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

최소 신장 트리를 만들기 위해 프림 알고리즘을 사용했다.

 

 

import sys
import heapq

input = sys.stdin.readline

V, E = map(int, input().split())

graph = [[] for i in range(V + 1)]

for _ in range(E):
    A, B, C = map(int, input().split())

    graph[A].append((C, (A, B)))
    graph[B].append((C, (B, A)))

count = 0
dp = [0] * (V + 1)

order = []

now_node = 1  # 시작 노드
result = 0  # 가중치 합

for edge in graph[1]:
    heapq.heappush(order, edge)

while count < V-1:

    # print(order)

    node = heapq.heappop(order)

    A = node[1][0]
    B = node[1][1]
    val = node[0]

    if dp[A] == 0 or dp[B] == 0:
        dp[A] = 1
        dp[B] = 1

        result += val

        now_node = B
        count += 1

        for edge in graph[now_node]:
            heapq.heappush(order, edge)

print(result)

 

 

방문한 노드가 들어있는 집합 A가 있다고 하자. 

(결국 모든 노드를 지나가기 때문에, 처음에는 시작 노드로 임의의 숫자를 넣어준다.)

 

집합 A의 노드와 연결된 간선 중에 가장 가중치가 작은 것을 찾아 지나가야 한다.  (그리디로 볼 수 있다.)

찾는 것은 최소힙을 이용했다.

 

N개의 노드를 방문하는 것은 간선 N-1개를 지나가는 것과 같다.

지나간 간선이 N-1개가 되면 알고리즘은 끝나고 최소 신장 트리가 완성된다.

 

 

 

처음 문제를 풀 때, 집합 A에 있는 노드와 연결된 모든 간선을 고려하지 않았다.

그래서 아래와 같은 반례에서 틀렸다.

 

3 3

1 2 2

1 3 3

2 3  9999

 

현재 노드[now_node]와 연결된 간선에서 길을 찾아야 한다는 생각 때문에

1->2->3이 돼서 결과가 10001이 나오게 된다.

 

1->2

1->3

이 될 수 있는데 말이다.

 

프림 알고리즘을 제대로 이해하고 푸니까 맞았다.