[알고리즘] 최소 스패닝 트리

난이도

Gold 4

문제

그래프가 주어졌을 때, 그 그래프의 최소 스패닝 트리를 구하는 프로그램을 작성하시오.

최소 스패닝 트리는, 주어진 그래프의 모든 정점들을 연결하는 부분 그래프 중에서 그 가중치의 합이 최소인 트리를 말한다.

입력

첫째 줄에 정점의 개수 V(1 ≤ V ≤ 10,000)와 간선의 개수 E(1 ≤ E ≤ 100,000)가 주어진다. 다음 E개의 줄에는 각 간선에 대한 정보를 나타내는 세 정수 A, B, C가 주어진다. 이는 A번 정점과 B번 정점이 가중치 C인 간선으로 연결되어 있다는 의미이다. C는 음수일 수도 있으며, 절댓값이 1,000,000을 넘지 않는다.

그래프의 정점은 1번부터 V번까지 번호가 매겨져 있고, 임의의 두 정점 사이에 경로가 있다. 최소 스패닝 트리의 가중치가 -2,147,483,648보다 크거나 같고, 2,147,483,647보다 작거나 같은 데이터만 입력으로 주어진다.

출력

첫째 줄에 최소 스패닝 트리의 가중치를 출력한다.

예제 입력

3 3
1 2 1
2 3 2
1 3 3

예제 출력

3

해설 및 후기

크루스칼 알고리즘의 구현이다. 프로그래머스에서 한번 해봐서 어렵진 않았지만, 제한이 조금 더 빡빡하다. 그래서 union-find 부분을 좀더 깔끔하게 해결하는 방법을 알게 되었다.

크루스칼 알고리즘의 간단한 설명은 다음과 같다. 그리디 알고리즘의 응용이다. 우선 리스트에 노드 a와 b, 그리고 이를 연결하는데 필요한 가중치를 저장한다. 그 리스트를 가중치가 낮은 순서대로 정렬한다. 다음으로 정렬된 리스트 순서대로 간선을 선택하는데, 노드의 순환이 발생하지 않는 선에서 간선을 선택한다.

노드의 순환이 발생하는 지 판단할 수 있는 방법은 union-find이다. 연결된 간선끼리 각 노드들을 union하여, 가지고 있는 집합이 서로 같으면 그 노드를 연결했을 때 순환이 발생한다. 집합을 모두 정의하여 이를 구현하기에는 시간복잡도와 공간복잡도 문제가 발생하므로, 각 집합의 root(번호가 가장 작은 노드)를 비교해 이를 구현한다.

제출 코드

import sys

v, e = map(int,input().split())
tree = []
for i in range(e):
    a,b,value = map(int, sys.stdin.readline().rstrip().split())
    tree.append([a,b,value])

parents = [i for i in range(v+1)]

def union(a,b):
    aRoot = find_root(a)
    bRoot = find_root(b)
    if(aRoot < bRoot):
        parents[bRoot] = aRoot
    else:
        parents[aRoot] = bRoot

def find_root(x):
    if(parents[x] != x):
        return find_root(parents[x])
    else:
        return x

tree.sort(key = lambda x:x[2])
cnt = 0
nCnt = 0
for i in range(e):
    if(nCnt == v-1):
        break
    a,b,value = tree[i]
    if(find_root(a) != find_root(b)):
        union(a,b)
        cnt += value
        nCnt += 1

print(cnt)