[Python] 백준 1197번 - 최소 스패닝 트리

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

1197번: 최소 스패닝 트리

 

처음에 Prim으로 풀었는데 예제만 맞고 또 일반화에 실패했다. ㅋㅋ

찾아보니까 Prim은 우선순위큐를 쓰더라고.. 복잡하길래 그냥 Kruskal로 풀었다.

유니온 파인드 공부해놨더니 별로 안 어려워서 풀만 했다.

 

import sys
input = sys.stdin.readline
def find(a):
    if a == parent[a]:
        return a
    parent[a] = find(parent[a])
    return parent[a]

def union(a, b):
    a = find(a)
    b = find(b)
    if a > b:
        parent[a] = b
    else:
        parent[b] = a

v, e = map(int, input().split())
edge = []
for i in range(e):
    a, b, c = map(int, input().split())
    edge.append([c, a, b]) #c: 가중치

edge.sort(key=lambda x:x[0])
parent = [i for i in range(v+1)]
count = 0 #노드 개수 체크
MST = 0
for i in range(e):
    weight = edge[i][0]
    start = edge[i][1]
    end = edge[i][2]
    if find(start) != find(end):
        union(start, end)
        MST += weight
        count += 1
    if count == v-1:
        break
print(MST)