문제 출처

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

본 문제는 노드들과 노드 간 가중치가 주어졌을 때, 최소 스패닝 트리의 총 가중치 합을 구하는 문제이다.

본 문제를 풀기 위해서 Prim Algorithm을 사용하였다.

풀이

다음과 같은 그래프가 주어졌다고 가정하자. 그래프는 위키피디아 한글 문서의 예제를 사용하였다.

7 11
1 2 7
1 4 5
4 2 9
4 6 6
4 5 15
2 3 8
2 5 7
6 5 8
6 7 11
3 5 5
5 7 9

우선 임의의 노드 하나를 선택한다. 여기서는 1번 노드를 선택하겠다.

1번 노드에서는 2번 노드와 4번 노드에 방문할 수 있다. 이 중 최소값은 4번 노드에 방문하는 5가 되므로 해당 노드를 선택해준다.

4번 노드에 방문했으므로 가능한 2번 노드, 5번 노드, 6번 노드에 추가적으로 방문할 수 있게 되었다. 이 중 최소값은 6번 노드에 방문하는 가중치 6이 되므로 해당 노드를 선택해준다.

6번 노드에 방문했으므로 5번 노드, 7번 노드에 추가적으로 방문할 수 있게 되었다. 이중 최소값은 2번 노드에 방문하는 가중치 7이 되므로 해당 노드를 선택해준다.

2번 노드에 방문했으므로 5번 노드, 3번 노드에 추가적으로 방문할 수 있게 되었다. 이중 최소값은 5번 노드에 방문하는 가중치 7이므로 해당 노드를 선택해준다.

5번 노드에 방문했으므로 3번 노드, 7번 노드에 추가적으로 방문할 수 있게 되었다. 이 중 최소값은 3번 노드에 방문하는 가중치 5이므로 해당 노드를 선택해준다.

마지막으로 7번 노드가 남았으므로 7번 노드에 방문해준다.
여기서 5번 노드에서 7번 노드로 가는 가중치는 9이지만, 2에서 3으로 가는 가중치 8, 5에서 6으로 가는 가중치 8이 남아있다. 그러나 2,3,5,6번 노드가 모두 방문이 되어있기 때문에 추가적으로 방문하지 않는다.

즉 프림 알고리즘을 정리하자면 다음과 같다.

1. 임의의 노드를 선택한다.
2. 해당 노드에서 갈 수 있는 노드 중, 가장 가중치가 적은 노드를 선택하여 방문한다.
3. 방문한 노드로부터 갈 수 있는 노드 중, 가장 가중치가 적은 노드를 선택하여 방문하는 것을 반복한다.
   • 만약 가장 가중치가 적은 노드가 이미 방문되었다면 무시하고 다음으로 가중치가 적은 노드를 선택한다.

코드

import sys 
import heapq as h

input = sys.stdin.readline

V, E = map(int, input().split())
dic = {} # start : (end, weight)

for _ in range(E): # edge의 개수만큼
    start, end, weight = map(int, input().split())
    if start not in dic.keys():
        dic[start] = []
    if end not in dic.keys():
        dic[end] = []

    dic[start].append((end, weight))
    dic[end].append((start ,weight))

#print(dic)
#print(visited)

start_node = 1 # 4에서부터 시작
V_tbs = start_node # to be searched

next_candidates = [] # 갈 수 있는 간선 (weight, start, end) 형태로 저장
h.heapify(next_candidates)

result = 0
num_searched = 1 # 탐색된 node의 개수 (start_node 때문에 1)
visited = [False] * (V+1)
visited[start_node] = True

while True:

    for v in dic[V_tbs]:
        h.heappush(next_candidates, (v[-1], V_tbs, v[0])) # heap에 weight, start, end 저장

    # 최소 weight
    #print(next_candidates)
    while True:
        min_ = h.heappop(next_candidates)
        if visited[min_[-1]] == False:
            visited[min_[-1]] = True
            break

    #print(min_)

    result += min_[0]
    V_tbs = min_[-1]
    num_searched += 1

    if num_searched >= V:
        break

print(result)