[snippet] kruskal.py

백준 1197번. 최소 스패닝 트리에 대한 풀이를 Kruskal's Algorithm의 스니펫 느낌으로 작성.

• MST(Minimum Spanning Tree; 최소 신장 트리) 를 구할 때 사용.

import sys
input = lambda: sys.stdin.readline().rstrip()

V, E = map(int, input().split())
parent = [i for i in range(V+1)]

edges = []
for _ in range(E):
    A, B, cost = map(int, input().split())
    edges.append((A, B, cost))

# union-find
def get_parent(x):
    if parent[x] == x:
        return x
    parent[x] = get_parent(parent[x]) # get_parent 거슬러 올라가면서 parent[x] 값도 갱신
    return parent[x]

def union_parent(a, b):
    a = get_parent(a)
    b = get_parent(b)
    
    if a < b: # 작은 쪽을 parent로 합의하자
        parent[b] = a
    else:
        parent[a] = b        

def same_parent(a, b):
    return get_parent(a) == get_parent(b)


################################

ans = 0
for a, b, cost in sorted(edges, key=lambda x: x[2]):
    # cost가 작은 edge부터 순회하면서, 같은 부모를 공유하지 않을 때(사이클 없을 때)만 확정
    if not same_parent(a, b):
        union_parent(a, b)
        ans += cost
print(ans)

################################