Problem Link
https://www.hackerrank.com/challenges/primsmstsub/problem?isFullScreen=true
프림 알고리즘(Prim's algorithm)을 이용한 최소 신장 트리(Minimum spanning tree) 탐색 문제
1. 프림 알고리즘(Prim's algorithm)
그림 출처: Wikipedia
from collections import defaultdict
import heapq
def prims(n, edges, start):
visited = [0] * (n + 1)
graph = defaultdict(list)
for u, v, w in edges:
graph[u].append([w, u, v])
graph[v].append([w, v, u])
heapq.heapify(graph[start])
visited[start] = 1
total_weight = 0
while graph[start]:
w, u, v = heapq.heappop(graph[start])
if visited[v] == 0:
visited[v] = 1
total_weight += w
for edge in graph[v]:
if visited[edge[2]] == 0:
heapq.heappush(graph[start], edge)
return total_weight
📌 코드 구현 설명
- 간선 정보가 저장된
edges
를 이용하여graph
를 만든다.heapq
라이브러리를 사용하여graph
를 최소 힙(Min heap) 형태로 만든다.
- 가중치가 가장 작은 간선(edge)을 우선적으로 선택해야 하므로 최소 힙을 사용한다.
- 모든 노드를 방문할 때 까지
graph
를 탐색한다.
visited
에 방문 여부를 기록하여 재방문하지 않도록 한다.- 노드를 방문할 때 마다
total_weight
에 가중치를 더하고, 반복문이 종료되면return
하여 정답을 구한다.