practivceAlgorithm/swexpertacademy

[SWEA][Python] 5249 최소 신장 트리

findTheValue 2021. 10. 10. 03:23

기본적인 MST 문제로 간선이 주어졌기 때문에 크루스칼 알고리즘을 사용해 풀이했습니다.

 

from heapq import heappop, heappush


def find(x):
    if parents[x] == x:
        return x
    parents[x] = find(parents[x])
    return parents[x]


def union(a, b):
    a = find(a)
    b = find(b)

    if a > b:
        a, b = b, a
    parents[b] = a


for test in range(1, int(input()) + 1):
    V, E = map(int, input().split())
    parents = {i: i for i in range(V + 1)}
    edges = []
    for __ in range(E):
        a, b, w = map(int, input().split())
        heappush(edges, (w, a, b))

    answer = 0
    while edges:
        w, a, b = heappop(edges)
        if find(a) != find(b):
            union(a, b)
            answer += w
    print(f'#{test} {answer}')