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

이 문제는 문제 제목에서부터 MST 문제라는 것을 알려주고 있다. 당연히 나는 MST 알고리즘 중 내가 알고 있는 Prim알고리즘을 통해 풀었고, 아쉽게도 시간 통과를 하지 못하였다.

잘 모르는 크루스칼 알고리즘을 공부해봐야 하나 생각하며 백준 질문게시판을 둘러보았는데, 프림과 힙을 사용해서 풀었다는 사람들이 많아서 조금 충격을 받았다.

나의 다른 MST 풀이를 보면 알겠지만 나는 프림 알고리즘은 세 가지 구조를 가지고 문제를 해결한다.
1. 현재 내가 갈 수 있는 곳 중 최소 비용으로 갈 수 있는 곳을 찾는다
2. 최소로 갈 수 있는 곳을 등록한다.
3. 새로 등록한 곳에서 또 더 갈 수 있는 곳을 찾는다.
이렇게 할 때, 힙을 어디서 사용하지라는 아이디어가 쉽게 떠오르지 않았다. 하지만, 결국 프림 알고리즘이 다익스트라와 유사한 알고리즘이라는 말에서 힌트를 얻었다. 

아이디어는 이렇다.
1. 현재 내가 갈 수 있는 곳(방문 안 한 곳)을 힙 자료구조에 넣는다(비용 기준)
2. 비용이 작은 곳부터 꺼내어 (Heappop) 또 갈 수 있는 모든 곳을 힙에 넣는다.(위 과정에 3번 과정)

문장으로도 한 단계도 줄 수 있고, 최솟값을 for문이 아닌 힙 자료구조를 통해 바로 찾을 수 있어 공간, 시간 메모리 상 위에 과정보다 확실히 단축시킬 수 있었다. 위에 과정으로 프림 알고리즘을 구현할때는 다익스트라랑 유사하다는 것을 느끼지 못하였는데 이 코드를 구현하면서 정말 비슷하다고 느낄 수 있었다.

코드로 나타내면 아래와 같다.

from sys import stdin
from heapq import heappop, heappush

input = stdin.readline

# Prim 알고리즘
def prim():
    global total
    V[1] = 1
    Q = []
    for c in route[1]:
        heappush(Q,c)


    while Q:
        w, idx = heappop(Q)
       if V[idx]:
            continue
        total += w
        V[idx] = 1
        for c in route[idx]:
            if not V[c[1]]:
                heappush(Q,c)

# 1. 변수저장
N, E = map(int,input().split())
route = [[] for _ in range(1+N)] # 간선 그래프
V = [0] * (1 + N) # 방문기록 그래프

for _ in range(E):
    p, c, w = map(int,input().split())
    route[p].append([w,c]) # 힙에 편하기 넣기 위해 가중치가 앞에 저장
    route[c].append([w,p])

total = 0
prim()
print(total)

프림알고리즘을 풀 때 for문이 세 개씩 있었던 나라서 항상 시간 복잡도를 걱정했는데 힙을 통해 이를 줄일 수 있는 방법을 배운 것 같아 뿌듯한 문제였다.

(그래도 크루스칼 알고리즘도 공부를 해둬야 할 것 같다..)