[SWEA][D4][#05249] 최소 신장 트리

출처

• SW Expert Academy

학습용 포스트입니다. 본 포스트가 문제가 될 시 수정 또는 삭제하겠습니다.

문제

그래프에서 사이클을 제거하고 모든 노드를 포함하는 트리를 구성할 때, 가중치의 합이 최소가 되도록 만든 경우를 최소신장트리라고 한다.

0번부터 V번까지의 노드와 E개의 간선을 가진 그래프 정보가 주어질 때, 이 그래프로부터 최소신장트리를 구성하는 간선의 가중치를 모두 더해 출력하는 프로그램을 만드시오.

입력

첫 줄에 테스트 케이스의 개수 T가 주어지고, 테스트 케이스 별로 첫 줄에 마지막 노드번호 V와 간선의 개수 E가 주어진다.

다음 줄부터 E개의 줄에 걸쳐 간선의 양 끝 노드 n1, n2, 가중치 w가 차례로 주어진다.

1<=T<=50, 1<=V<=1000, 1<=w<=10, 1<=E<=1000000

출력

각 줄마다 “#T” (T는 테스트 케이스 번호)를 출력한 뒤, 답을 출력한다.

예제

입력

3
2 3
0 1 1
0 2 1
1 2 6
4 7
0 1 9
0 2 3
0 3 7
1 4 2
2 3 8
2 4 1
3 4 8
4 6
0 1 10
0 2 7
1 4 2
2 3 10
2 4 3
3 4 10

출력

#1 2
#2 13
#3 22

My Sol

def find_set(n):
    if p[n] != n:
        p[n] = find_set(p[n])
    return p[n]


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

    cnt = 0
    ans = 0
    for edge in edges:
        s, e, w = edge
        a = find_set(s)
        b = find_set(e)
        if a==b: continue
        p[a] = b
        cnt += 1
        ans += w
        if cnt == V: break

    print(f'#{tc} {ans}')

크루스칼 알고리즘을 이용해 최단 신장 트리(MST)를 구하고, 가중치들의 합의 최솟값을 ans에 담아 출력하면 되는 문제였다.

크루스칼 알고리즘의 경우 입력을 가중치와 함께 모두 담아 가중치를 기준으로 오름차순 정렬해 각각의 간선을 선택하는데, 서로소 집합을 이용해 루트 노드의 번호를 반환하는 find_set 함수를 이용해 두 노드의 관계를 정의하고 가공한다. 두 노드가 같은 집합 내에 있다면 간선의 가중치를 고려해줄 필요가 없으므로 넘어가고, 다른 집합이라면 루트를 같게 만들어 준 뒤, 해당 간선의 가중치를 더해주는 것이다.

이 간선의 개수의 합이 V-1개, 즉 선택된 노드가 V개라면 종료한다.

결과

PASS