[Baekjoon 1197 / python / 골드4] 최소 스패닝 트리

import heapq

import sys

input = sys.stdin.readline

N, E = map(int, input().split())  # 노드, 간선

graph = [[] for _ in range(N + 1)]  # 그래프 입력받기

for _ in range(E):

    a, b, c = map(int, input().split())

    graph[a].append((c, b))

    graph[b].append((c, a))

visited = [False] * (N + 1)  # 방문체크 리스트

def prim(graph, start):  # 그래프, 시작노드 입력받는다, 프림 알고리즘 사용

    visited[start] = True  # 시작노드 방문체크

    distance = 0  # 거리(가중치)를 계산할 변수

    edge = [(cost, start, end) for cost, end in graph[start]]  # 시작노드의 인접노드에 대한 정보를 미리 리스트에 넣는다

    heapq.heapify(edge)  # 리스트를 우선순위 큐(최소힙)로 변경한다.

 

    while edge:  # 큐가 빌 때까지 진행( 모든 노드를 방문할 때까지 )

        cost, start, end = heapq.heappop(edge)  # 큐에서 가중치(거리) 가장 작은 값을 뺀다

        if visited[end] == False:  # 아직 안 간 곳이면

            visited[end] = True  # 방문체크

            distance += cost  # 거리 추가

            for nextCost, nextEnd in graph[end]:  # 방금 추가한 노드의 인접 노드에 대하여 반복

                if visited[nextEnd] == False:  # 안간 곳 있으면 큐에 추가하기

                    heapq.heappush(edge, (nextCost, end, nextEnd))

    return distance  # 계산된 최단거리 리턴

print(prim(graph, 1))

 

최소 스패닝 트리를 구하는 문제이다. 크루스칼 알고리즘과 프림 알고리즘을 사용할 수 있는데, 난 프림 알고리즘을 사용했다. 크루스칼도 알면 좋긴 한데 먼저 공부한 게 프림 알고리즘이기 때문이다.

 

heapq를 통한 우선순위 큐를 사용하며, 이것은 최소 힙이기 때문에 항상 가중치가 가장 적은 값이 pop되게 된다.

프림의 동작원리는 이렇다.

1. 임의로 시작 노드를 선택하고 T에 넣는다 ( T는 최소스패닝트리를 담을 변수 )

2. T에 있는 모든 노드의 인접 노드에 대하여, 아직 방문하지 않았고 가중치가 가장 적은 노드를 선택해 T에 넣는다.

3. 2번을 반복한다.

 

원리 자체는 엄청나게 복잡하진 않다. 또한 다익스트라 알고리즘과 비슷하다. 근데 최근에 이런저런 알고리즘을 단기간에 워낙 많이 공부해서 뭐가 뭔지 굉장히 헷갈려 프림 코드에 익숙해지는데에 애먹었다. 그래도 한번 이해하고 코드를 짜 놓으면 다음에 풀 땐 훨씬 수월하다.