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

문제

SW Expert Academy

풀이

백준 1197번, SWEA 3124번 문제와 같은 동일한 문제다. 차이점은 이번 문제는 Python으로만 제출이 가능하다는 것이다.

root = []

def find(x):
    if x == root[x]:
        return x
    else:
        root[x] = find(root[x])
        return root[x]

def union(A, B):
    rootA = find(A)
    rootB = find(B)
    if rootA == rootB:
        return False
    root[rootA] = rootB;
    return True

T = int(input())
for t in range(1, T+1):
    V, E = map(int, input().split())
    edges = list()
    for _ in range(E):
        n1, n2, w = map(int, input().split())
        edges.append([n1, n2, w])
    edges.sort(key=lambda x: x[2])

    root = list(i for i in range(V+1))

    answer = 0
    cnt = 0
    for A, B, W in edges:
        if union(A, B):
            answer += W
            cnt += 1
            if cnt == V:
                break
    print(f"#{t} {answer}")