[BOJ] 1368 - 물대기 (Python)

 

 

아이디어

 

크루스칼 알고리즘으로 최소 스패닝 트리를 구현하여 최소 비용을 구한다.

 

 

풀이

 

가상의 0번째 논을 설정해 우물을 팔 때 드는 비용을 0번 논과 연결하는 데 드는 비용으로 저장하여 푼다.

 

import sys

input = sys.stdin.readline

INF = int(1e8)

def solution():

    def find(n):
        temp = n
        while n != group[n]:
            n = group[n]
        group[temp] = n
        return n

    def union(n1, n2):
        n1 = find(n1)
        n2 = find(n2)
        if n1 > n2:
            n1, n2 = n2, n1
        group[n2] = n1

    N = int(input())
    group = [i for i in range(N+1)]
    edges = []
    for i in range(1, N+1):
        W = int(input())
        edges.append((W, 0, i))
    for i in range(1, N+1):
        P = tuple(map(int, input().split()))
        for j in range(i+1, N+1):
            edges.append((P[j-1], i, j))
    edges.sort()
    res = cnt = 0
    for p, i, j in edges:
        if find(i) != find(j):
            union(i, j)
            res += p
            cnt += 1
            if cnt >= N:
                break
    print(res)

solution()