[Python] 프로그래머스 / 섬 연결하기

문제 링크

 

최소 스패닝 트리

문제를 읽자마자 대놓고 "최소 스패닝 트리로 풀어!" 라고 말하고 있다.

최소 스패닝 트리(MST) 문제는 Prim 알고리즘과 Kruskal 알고리즘으로 풀 수 있다.

• Prim 알고리즘 : 임의의 정점에서 시작해서 가장 가까운 이웃 정점을 고르고, 다시 그 정점에서 가장 가까운 정점을 고르기를 반복하는 것
• Kruskal 알고리즘 : 간선 중 짧은 것부터 고르는 것. 다만, Cycle이 형성되면 안 됨.

Cycle이 형성되면 안 된다는 문제점 때문에 (이게 생각보다 코드 구현이 빡세서), 보통은 Prim 알고리즘이 선호된다. 여기에서도 Prim 방법을 쓸 것이다.

비슷한 백준 문제를 Kruskal 알고리즘으로 푼 포스트는 다음 링크에 있다:

https://loewen.tistory.com

 

정답 코드

import heapq
from collections import defaultdict

def solution(n, costs):
    world = defaultdict(list)
    
    for s, e, cost in costs:
        world[s].append([cost, e])
        world[e].append([cost, s])
    
    total_cost = 0
    
    seen = [False] * n
    queue = [[0, 0]]  # [비용, 노드]
    
    while queue:
        cost, node = heapq.heappop(queue)
        
        if seen[node]:
            continue
        
        seen[node] = True
        total_cost += cost
        
        for next_cost, nei in world[node]:
            if not seen[nei]:
                heapq.heappush(queue, [next_cost, nei])
    
    return total_cost