백준 1368 [물대기] 파이썬

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

 

# 백준 1368 물대기
# 문제 요약: 
# 각 논에 직접 우물을 팔거나, 다른 논에서 물을 끌어올 수 있다.
# 모든 논이 물을 공급받을 때 최소 비용을 구하는 문제.
# 풀이 아이디어: 
# "우물"을 가상의 0번 노드로 두고, 우물 ↔ 각 논의 비용(dig_costs)을 간선으로 생각.
# 그리고 논 ↔ 논 간의 물길 비용(recv_costs)도 간선으로 생각해서 MST를 구한다.
# 여기서는 Prim 알고리즘을 사용.

import sys
import heapq

input = sys.stdin.readline
N = int(input())  # 논의 개수

dig_costs = []    # dig_costs[i] = i번 논에 우물을 직접 파는 비용
recv_costs = []   # recv_costs[i][j] = i번 논에서 j번 논으로 물을 끌어오는 비용
visited = [False for _ in range(N + 1)]  # 각 논의 방문 여부 (0은 가상의 우물 노드)

# 각 논에 우물 파는 비용 입력
for _ in range(N):
    dig_costs.append(int(input()))

# 논 사이의 물길 비용 입력
for _ in range(N):
    recv_costs.append(list(map(int, input().split())))

# 우물을 팔 수 있는 모든 선택지를 우선순위 큐(힙)에 넣음
# (비용, 논 번호)
# → 이는 "가상의 우물 노드(0)에서 각 논으로 가는 간선"을 추가하는 효과
q = []
for i in range(N):
    heapq.heappush(q, (dig_costs[i], i + 1))

cnt = 0  # 물 공급된 논 개수
ans = 0  # 최소 비용 합

# Prim 알고리즘 시작
while q and cnt < N:
    cost, node = heapq.heappop(q)
    
    # 이미 물이 공급된 논이면 스킵
    if visited[node]:
        continue

    # 새로운 논에 물 공급
    visited[node] = True
    ans += cost
    cnt += 1

    # 현재 논에서 다른 논으로 물길을 연결하는 간선 후보들을 힙에 추가
    for nxt in range(1, N + 1):
        if not visited[nxt]:
            heapq.heappush(q, (recv_costs[node - 1][nxt - 1], nxt))

# 모든 논에 물 공급이 끝났으므로 최소 비용 출력
print(ans)

 

자면서 푸는법이 얼핏 떠올랐는데, 정리가 안됐다...

결국 GPT의 도움을 받았는데, 풀이를 보니 너무 허무했다.

그냥 우물도 heapq에 넣으면 그만이었다. (그래야 Prim 알고리즘이 정확히 동작)