[Python] 백준 / gold / 16562 : 친구비

문제 링크 : https://www.acmicpc.net/problem/16562

알고리즘 유형 : Union-Find

(생각의 흐름)

1. 친구의 친구는 친구다 -> 친구 집합을 만들어야겠다.
   union-find 를 사용한다.

2. 모든 친구를 사귀기 위해 필요한 '최소 비용'만 구하면 되기 때문에, 친구 set 을 구하고, 그 set 에 속한 친구들 중 비용이 가장 적은 친구 한 명만 사귀면 된다.

(기억할 것)
나중에 부모(루트)를 get 해야 할땐, 만들어뒀던 find 함수를 사용한다.
경로 압축을 사용했어도, parent[x] 엔 반드시 x 의 루트노드가 담긴다는 보장이 없다.

정답 코드

import sys
N, M, K = map(int, sys.stdin.readline().split())
cost = [0] + list(map(int, sys.stdin.readline().split()))
parent = [x for x in range(N+1)]


def findParent(x):
    if parent[x] == x:
        return x
    else:
        parent[x] = findParent(parent[x])
        return parent[x]


def union(x, y):
    px = findParent(x)
    py = findParent(y)

    if px != py:
        if px < py:
            parent[py] = px
        else:
            parent[px] = py


for _ in range(M):
    v, w = map(int, sys.stdin.readline().split())
    union(v, w)

parents = set()
for p in parent[1:]:
    parents.add(p)

parent_cost = dict()
for c in range(1, N+1):
    pc = findParent(c)
    if pc not in parent_cost:
        parent_cost[pc] = cost[c]
    else:
        parent_cost[pc] = min(parent_cost[pc], cost[c])

total_cost = 0
for key in parent_cost:
    total_cost += parent_cost[key]

if total_cost > K:
    print("Oh no")
else:
    print(total_cost)