- 우선순위 큐(heap)에서 비용이 가장 작은 간선을 꺼내 방문하지 않은 노드를 선택
- 선택한 노드를 MST(Minimum Spanning Tree, 최소 신장 트리)에 포함시키고, 그 노드에서 나가는 간선들을 큐에 추가
- 모든 노드를 방문할 때까지 반복해 MST의 총 비용을 계산하기
from collections import defaultdict
import heapq
def prim_mst(graph, start, n):
hp = []
visited = set()
heapq.heappush(hp, (0, start))
total_cost = 0
# heap에서 가장 최소 cost를 가진 엣지가 가장 먼저 팝 됨 -> 최소 비용으로 모든 노드 연결 가능
while hp:
cost, node = heapq.heappop(hp)
# heap에서 어떤 엣지부터 팝될지 모르므로 팝한 후에도 다시 방문여부 체크 필수!
if node in visited:
continue
total_cost += cost
visited.add(node)
for weight, child in graph[node]:
if child not in visited:
heapq.heappush(hp, (weight, child))
return total_cost
def solution(n, costs):
answer = 0
graph = defaultdict(list)
for i in range(len(costs)):
a, b, cost = costs[i][0], costs[i][1], costs[i][2]
graph[a].append((cost, b)) # cost 먼저 저장
graph[b].append((cost, a))
return prim_mst(graph, 0, n)
