[백준] 16562 친구비 python

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

16562번: 친구비

Union find 알고리즘을 사용하는 문제이다.

두 노드를 받아서 연결을 시키고, 친구가 될수 있는 최소 비용도 같이 갱신한다.

모든 노드를 받은 후 전체적으로 각 노드의 root값을 한번 더 갱신하고, 각 트리의 루트 비용를 더하면 된다.

# import sys; input = sys.stdin.readline
from collections import Counter
def find(x, root, rank):
    if root[x] == x:
        return x
    root[x] = find(root[x], root, rank)
    return root[x]

def union(x, y, root, rank, friend_cost):
    x = find(x, root, rank)
    y = find(y, root, rank)
    if x == y:
        return

    if rank[x] < rank[y]:
        root[x] = y
    else:
        root[y] = x
    if friend_cost[x] < friend_cost[y]:
        friend_cost[y] = friend_cost[x]
    else:
        friend_cost[x] = friend_cost[y]

    if rank[x] == rank[y]:
        rank[x] += 1
    return


def main():
    N, M, K = map(int, input().split())
    A = [0]
    A.extend(map(int, input().split()))
    root = list(range(N + 1))
    rank = [0 for _ in range(N + 1)]
    for _ in range(M):
        v, w = map(int, input().split()) # friend
        union(v, w, root, rank, A)

    # root renew
    for i in range(1, N + 1):
        root[i] = find(i, root, rank)

    # total root node
    total = 0
    for k in Counter(root).keys():
        total += A[k]
    print(total if total <= K else "Oh no")



if __name__ == '__main__':
    main()

아직 union find는 이해가 가지만, 안보고 짜기가 어려워서 한번더 풀어봐야할듯 하다.