[백준] 16562 - 친구비 Python

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

16562번: 친구비

친구가 없는 준석이가 "친구의 친구는 친구다" 라는 개념을 활용해 가장 적은 비용으로 모든 사람과 친구가 되는 최소 비용을 구하는 문제다.

 

친구의 친구임을 활용하기 때문에, 노드를 친구라는 집합 개념으로 묶어주는 전형적인 union-find 문제임을 알 수 있다.

- 친구 집합에 포함되어 있는지 find 를 이용한다.

- 친구끼리는 union 을 이용해 같은 집합으로 합친다.

 

최소 비용으로 모든 친구와 친구가 되어야 하므로, union 에서 노드 번호가 작은쪽으로 합치지 않고 비용이 저렴한 쪽으로 합치는것이 핵심이다. 비용이 적은 쪽으로 합치게 되면, 노드별로 find 할 때 노드가 속한 집합에서 가장 비용이 적은 노드의 번호를 반환한다. 결과적으로 find 의 결과가 자기 자신의 번호와 같다면, 집합에서 가장 비용이 적은 노드가 된다.  

 

코드는 다음과 같다.

from sys import stdin


def find(x):
    if x != parent[x]:
        parent[x] = find(parent[x])
    return parent[x]


def union(a, b):
    a = find(a)
    b = find(b)
    if c[a] < c[b]:
        parent[b] = a
    else:
        parent[a] = b


N, M, k = map(int, stdin.readline().split())
parent = [x for x in range(N+1)]
c = [0] + list(map(int, stdin.readline().split()))
for _ in range(M):
    a, b = map(int, stdin.readline().split())
    union(a, b)

ans = 0
for i in range(1, N+1):
    if i == find(i):
        ans += c[i]

print(ans if k >= ans else 'Oh no')