[Python] 백준 1197번 최소 스패닝 트리

1197번: 최소 스패닝 트리

문제

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

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

입력

첫째 줄에 정점의 개수 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보다 작거나 같은 데이터만 입력으로 주어진다.

출력

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

코드

# 크루스칼 
import heapq

INF = 1000001
v, e = map(int, input().split())
edge = []
parent = [i for i in range(v+1)]

for _ in range(e):
    a,b,c = map(int, input().split())
    heapq.heappush(edge, (c, a, b))
       
def find_parent(parent, x):
    # 루트 노드가 아니라면 루트 노드 찾을때까지 재귀적으로 호출
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])
    return parent[x]

def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    
    if a == b:  # 루트 노드가 같으면 사이클 발생
        return False
    # 합치기
    if a < b:
        parent[a] = b   # 정점 번호가 더 작은 경우를 루트쪽으로
    else:
        parent[b] = a
    return True

ans = 0

while edge:
    cost, a, b = heapq.heappop(edge)
    if union_parent(parent, a, b):
        ans += cost
print(ans)

학교에서 자료구조를 수강할 때 배웠던 내용이다. 최소 스패닝 트리(MST, Minimum Spanning Tree)를 구하는데에는 세 가지 알고리즘을 사용할 수 있다. 

1. 크루스칼(Kruskal) 알고리즘
2. 프림(Prim) 알고리즘
3. 솔린(Solin) 알고리즘

이 중 솔린은 잘 사용하지 않는다. 보통 크루스칼과 프림을 많이 사용한다. 나는 이 문제에서는 크루스칼 알고리즘을 사용하여 풀이를 하였다. 

크루스칼 알고리즘은 그리디 알고리즘을 이용하여 MST를 만드는 방법인데 간단히 설명하자면 간선들의 가중치를 오름차순으로 정렬하여 순서대로 살펴보며 사이클을 형성하지 않는 간선들을 선택하여 트리를 만드는 방법이다.

1. 간선들의 정보를 최소힙에 입력받으며 가중치의 오름차순으로 정렬되도록 한다.
2. find_parent 함수는 해당 노드의 루트 노드를 찾아주는 함수.
3. union_parent 함수는 인자로 받은 a, b 노드를 합쳐주는 함수인데 두 노드의 루트 노드를 찾았을 때 두 노드의 루트 노드가 같다면 합쳤을 때 사이클이 발생한다는 말이므로 이 경우에는 합쳐주지 않고 False를 리턴하고 함수가 종료된다. 사이클이 발생하지 않는다면 두 노드는 합쳐지고 True를 리턴한다.
4. 최소힙에 저장된 간선 정보들을 하나씩 pop하며 사이클을 형성하지 않는다면 트리를 합쳐준다. 이때 트리가 합쳐진다면 가중치를 ans 변수에 더하여 문제의 값을 구해나간다.