[Python] boj no. 1368 : 물대기

Problem

https://www.acmicpc.net/problem/1368

Idea

정점 간 간선의 길이가 주어지고 최소 비용으로 모든 정점을 연결하는 경로를 구하라는 것을 통해 최소 신장 트리에 관한 문제라는 것을 바로 알았다.

prim 알고리즘을 사용하여 최소 신장 트리를 구하되, 시작 정점을 우물을 파는데 가장 적은 비용이 드는 정점으로 설정하고 시작 정점을 제외한 모든 정점에 대해 자기 자신을 도착 정점으로 하며 비용이 우물 파는 비용인 간선을 초기 간선 집합에 넣어 주었다. prim 알고리즘을 통해 현재 신장 트리 집합에 트리 집합과 가장 가깝거나 우물 파는 비용이 가장 낮은 정점 중 더 작은 것을 추가하여 집합이 확장될 것이다.

Code

import sys
import heapq

def prim():
    selected = [False for _ in range(N)]
    sum = 0
    
    while heap:
        weight, vertex = heapq.heappop(heap)
        if selected[vertex] == False:
            selected[vertex] = True
            sum += weight
            
            for i, w in enumerate(edge[vertex]):
                if selected[i] == False:
                    heapq.heappush(heap, (w, i))
                    
        if False not in selected:
            break          
            
    print(sum)

input = sys.stdin.readline

N = int(input())
heap = []
for i in range(N):
    heapq.heappush(heap, (int(input()), i))

edge = [[] for _ in range(N)]
for i in range(N):
    Wi = list(map(int, input().split()))
    edge[i] = Wi

prim()

Comment