백준 1368 <물대기> Python

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

---

최소 스패닝 트리 문제입니다.

 

그냥 크루스칼 알고리즘을 돌리면 되는데, 우물을 파는 비용을 고려해주어야 하기 때문에 한 가지 수정이 필요합니다. 인덱스가 N+1 인 가상의 노드를 생성한 뒤 이 노드는 '우물을 파는 가격' 에 대한 노드로 설정해줍니다. 즉 크루스칼 알고리즘을 이 N+1 노드를 포함해서 돌리게 되면 다른 노드와 연결하는 비용과 우물을 파는 비용들을 모두 합쳐 비교해 최소 비용을 계산할 수 있습니다.

---

import sys

input = sys.stdin.readline


## Union-Find

# Find
def Find(group_list, node):
    # 루트 노드가 아니면
    if group_list[node] != node:
        group_list[node] = Find(group_list, group_list[node])
    
    return group_list[node]


# Union
def Union(group_list, n1, n2):
    # 그룹
    g1 = Find(group_list, n1)
    g2 = Find(group_list, n2)
    # 같은 그룹이면
    if g1 == g2:
        # 병합하지 않음
        return False, group_list

    # 작은 쪽으로 병합
    if g1 < g2:
        group_list[g2] = g1
    
    else:
        group_list[g1] = g2

    return True, group_list


def solution(N, Ws, Es):
    # 최소비용
    min_cost = 0
    # 그룹
    group_list = list(range(N+1))
    # 엣지 정리
    Es = [[i, j, Es[i][j]] for i in range(N) for j in range(N) if i < j]
    # 우물을 파는 비용을 가상의 논에 연결하는 셈 치자
    for i in range(N):
        Es.append([i, N, Ws[i]])
    
    # 정렬
    Es.sort(key=lambda x:x[-1])
    # 순회하며 병합
    for a, b, c in Es:
        state, group_list = Union(group_list, a, b)
        # 병합 되었으면
        if state:
            # 최소비용 추가
            min_cost += c
    
    print(min_cost)


# 입력
N = int(input().strip())
Ws = [int(input().strip()) for _ in range(N)]
Es = [list(map(int, input().strip().split())) for _ in range(N)]

solution(N, Ws, Es)