프림 - 1368번: 물대기

jisu_log·2025년 9월 13일

알고리즘 문제풀이

목록 보기
101/105

  • 각 논에 우물 파는 비용을 “가상 정점에서 논으로 가는 간선”으로 초기 힙에 모두 삽입하기 (우물 비용 가장 싼 논에서 시작하는 걸로 하면 틀림! MST가 하나만 존재하는게 아니기 때문에)
  • 힙에서 가장 싼 비용을 꺼내 방문하지 않은 논을 선택, 전체 비용 누적
  • 선택된 논에서 다른 논으로 연결되는 파이프 비용을 힙에 추가하며 최소 비용 신장트리 만들기
    -> 우물 비용 vs 파이프 비용 중 자동으로 최소가 선택되도록(힙 이용) 만드는 것!
from collections import defaultdict
import heapq
import sys

input = sys.stdin.readline

N = int(input())
w_list = []
graph = defaultdict(list)

for i in range(N):
    W = int(input())
    w_list.append(W)

for i in range(N):
    line = list(map(int, input().split()))
    for j in range(N):
        if i == j:
            continue
        graph[i].append((j, line[j]))


def prim():
    total_fee = 0
    heap = []
    for i in range(N):
        heapq.heappush(heap, (w_list[i], i))
    visited = [False] * N

    while heap:
        cost, node = heapq.heappop(heap)

        # pop한 후에도 방문체크 다시 해줘야 함(heap은 push된 순서대로 pop되지 않으므로)
        if visited[node]:
            continue

        visited[node] = True
        total_fee += cost

        # 현재 방문한 노드의 이웃 간선들 푸시하기
        for c, w in graph[node]:
            if not visited[c]:
                heapq.heappush(heap, (w, c))

    return total_fee


print(prim())

0개의 댓글