문제
논에 물을 대는 방법은 두 가지가 있다. 우물을 파거나 다른 논에서 물을 끌어오거나. 모든 논에서 다른 논으로 물을 끌어올 수 있고 모든 논에서 우물을 팔 수 있을 때 모든 논에 물을 대는 최소비용을 구하시오.
풀이
아이디어 하나만 있으면 쉬운 문제가 된다. 우물을 파는 경우를 0번째 논으로 지정하고 처음부터 visited상태로 놓고 heapq에 우물을 파는 경우 cost를 모두 넣고 시작하면 된다.
이렇게 하면 단순 Prim's 알고리즘으로 풀린다.
import heapq
n = int(input())
board = [[0]]
for i in range(n):
cost = int(input())
board[0].append(cost)
for i in range(n):
board.append([board[0][i+1]] + list(map(int,input().split())))
visited = [False]*(n+1)
visited[0] = True
total_cost = 0
q = []
for i in range(1,n+1):
heapq.heappush(q,(board[0][i],i))
while q:
cost,node = heapq.heappop(q)
if visited[node]:
continue
total_cost += cost
visited[node] = True
for i in range(n+1):
if visited[i]:
continue
cost = board[i][node]
heapq.heappush(q,(cost,i))
print(total_cost)