[백준 1368] 물대기

1. 문제 설명

물대기

2. 문제 분석

크루스칼 알고리즘을 통해 최소 신장 트리를 구한다. 이때 i->j 간선 비용과 i->i 간선 비용을 비교해야 하는데, i->i 간선 비용을 0->i 간선으로 처리한다.

3. 나의 풀이

import sys
import heapq

n = int(sys.stdin.readline().rstrip())
pq = []
parents = [i for i in range(n+1)]
for i in range(1, n+1):
    heapq.heappush(pq, [int(sys.stdin.readline().rstrip()), 0, i])
    # i번에 직접 물을 대는 경우는 0번에서 i로 이어지는 간선으로 간주

for i in range(1, n+1):
    line = list(map(int, sys.stdin.readline().rstrip().split()))
    for j in range(1, n+1):
        if i == j: continue
        heapq.heappush(pq, [line[j-1], i, j])

def find(node):
    if parents[node] == node: return node
    else:
        parents[node] = find(parents[node])
        return parents[node]

def union(node1, node2):
    root1, root2 = find(node1), find(node2)
    if root1 == root2: return False
    else:
        parents[root2] = root1
        return True

total = 0
while pq:
    cur_cost, node1, node2 = heapq.heappop(pq)
    if union(node1, node2):
        total += cur_cost
    # 크루스칼 알고리즘

print(total)

좋은 웹페이지 즐겨찾기