[Python] 백준 16562 - 친구비 문제 풀이

Overview

BOJ 16562번 친구비 Python 문제 풀이
분류: Disjoint Set (분리집합), Union Find (유니언 파인드)

---

문제 페이지

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

---

풀이 코드

from sys import stdin


parents = []
charges = []


def find(x: int) -> int:
    if parents[x] == x:
        return x
    parents[x] = find(parents[x])
    return parents[x]


def union(x: int, y: int) -> None:
    x, y = find(x), find(y)
    # 친구비를 기준으로 부모 설정
    if charges[x] < charges[y]:
        parents[y] = x
    else:
        parents[x] = y


def main():
    def input():
        return stdin.readline().rstrip()
    global parents, charges

    n, m, k = map(int, input().split())
    parents = list(i for i in range(n + 1))
    charges = [0] + list(map(int, input().split()))
    for _ in range(m):
        v, w = map(int, input().split())
        union(v, w)

    friends = set()
    cost = 0
    # 친구들 중 parent인 친구의 요금 총합 구하기
    for i in range(1, n + 1):
        if find(i) not in friends:
            cost += charges[parents[i]]
            friends.add(parents[i])

    if cost > k:
        print("Oh no")
    else:
        print(cost)


if __name__ == "__main__":
    main()

유니언 파인드를 응용하여 풀이하였다.
union 함수에서 각 노드의 parent를 정할 때, 친구비(charge)가 작은 쪽을 parent로 설정하였다.
입력으로 들어온 친구들을 union하고, 전체 친구들 중에서 parent인 친구들의 친구비 합을 구한다.