[BOJ] 6497 - 전력난 (Python)

 

 

아이디어

 

크루스칼 알고리즘을 사용한 최소 신장 트리로 집을 모두 연결한 최소 비용을 전체 비용에서 뺀다.

 

 

풀이

 

최소 신장 트리를 구현하기 위해 집합을 찾는 find 함수와 두 집합을 합치는 union 함수를 정의한다.

집의 수는 M, 길의 수는 N으로 입력 받고, while문은 집의 수 M이 유효한 동안 반복된다.

길에 대한 정보를 거리 순으로 정렬한 road 리스트를 만들고 유니온 파인드로 두 집이 연결되지 않았을 때 연결한다.

road 리스트를 만들 때 전체 비용 합을 res에 저장해두고, 두 집이 연결될 때마다 res에서 해당 비용을 뺀다.

연결된 횟수(간선 개수) cnt가 M-1이 되면 for문을 종료하고 결과값 res를 출력한다.

 

import sys

input = sys.stdin.readline

def find(n):
    temp = n
    while group[n] != n:
        n = group[n]
    group[temp] = n
    return n

def union(n1, n2):
    group[find(n2)] = find(n1)

while 1:
    M, N = map(int, input().split())
    if not M:
        break
    road = []
    res = 0
    for i in range(N):
        x, y, z = map(int, input().split())
        road.append((x, y, z))
        res += z
    road.sort(key=lambda r: r[2])
    group = [i for i in range(N)]
    cnt = 0
    for x, y, z in road:
        if find(x) != find(y):
            union(x, y)
            res -= z
            cnt += 1
            if cnt >= M-1:
                break
    print(res)

 

 

리스트 컴프리헨션보다 for문에서 append하는 코드가 더 빨랐다.

후자의 경우 전체 비용 합을 for문에서 미리 구할 수 있어 이후 정답을 구하는 for문을 다 돌지 않고도 종료할 수 있었다.

# 1.
road = [list(map(int, input().split())) for _ in range(N)]

# 2. 더 빠른 코드
res = 0
for i in range(N):
    x, y, z = map(int, input().split())
    road.append((x, y, z))
    res += z

 

리스트는 mutable, 튜플은 immutable하므로 속도면에서 튜플이 더 유리했다.

# 1.
road.append([x, y, z])

# 2. 더 빠른 코드
road.append((x, y, z))

 

결과값을 문자열이 아닌 리스트로 받고 for문에서 슬라이싱하여 언패킹할 수도 있다.

큰 차이는 아니지만 메모리가 150KB 정도 더 사용되기에 리스트가 아닌 문자열로 받도록 했다.

# 메모리가 더 많이 사용되는 코드
res = [mid]
...
if not j%2:
    res.append(mid)
...
print(len(res))
for i in range(len(res)//10+1):
    print(*res[i*10:(i+1)*10])