백준 1197 [최소 스패닝 트리] 파이썬

https://www.acmicpc.net/problem/1197

 

스패닝 트리 개념은 알지만 실제 구현하려니 어렵다..

 

<실패 코드>

# 백준 1197 최소 스패닝 트리
import sys
sys.setrecursionlimit(10**6)

V,E=map(int,sys.stdin.readline().split())
visited=[False]*(V+1)
lines=[]
for _ in range(E):
    lines.append(list(map(int,sys.stdin.readline().split())))

lines.sort(key=lambda x: -x[2])
ans=0


cnt=0

while cnt<V:
    if lines:
        A,B,cost=lines.pop()   
        if visited[A] and visited[B]:
            continue
        elif not visited[A] and not visited[B]:
            visited[A]=True
            visited[B]=True
            ans+=cost
            cnt+=2
        else:
            visited[A]=True
            visited[B]=True
            ans+=cost
            cnt+=1

print(ans)

방문하지 않았다면 방문하여 간선을 연결하는 방식으로 하였는데, 생각해보니 여러개의 트리가 만들어지는 경우가 있을 수도 있다.

 

따라서 union-find 방식을 사용하여 루트 노드를 find 하고, 루트 노드가 다르다면(서로 다른 트리)이므로 연결해주는 방식으로 해결하는게 좋다.

 

<정답 코드>

# 백준 1197 최소 스패닝 트리 (크루스칼)
import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline

# 유니온 파인드 구현
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])  # 경로 압축
    return parent[x]

def union(x, y):
    x = find(x)
    y = find(y)
    if x != y:
        parent[y] = x
        return True
    return False

V, E = map(int, input().split())
edges = [tuple(map(int, input().split())) for _ in range(E)]
edges.sort(key=lambda x: x[2])  # 비용 기준 오름차순

parent = [i for i in range(V+1)]

ans = 0
cnt = 0

for a, b, cost in edges:
    if union(a, b):
        ans += cost
        cnt += 1
        if cnt == V - 1:  # MST는 V-1개의 간선으로 구성
            break

print(ans)