[Python] 16562 - 친구비 (골드4)

1. 문제 설명

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

16562번: 친구비

내가 이렇게 못 받은 친구비만 10억이 넘어..

 

---

2. 아이디어

https://jaeseo0519.tistory.com/61?category=1071535 

[Python] 1717 - 집합의 표현 (골드4)

 

이 문제를 풀어보면 union-find 알고리즘에 대해 쉽게 이해할 수 있다.

대신 이번에는 정점을 이동하는데 비용이 추가하고, 가장 비용이 적게 드는 경우를 찾아야 한다.

 

우선 처음에 친구 정보를 갱신해주는데 이때 준석이를 기준으로 생각해보자.

union은 정점의 합집합 정보 여부가 쌍방향으로 작용하지 않는다.

A가 B 정점으로 연결되어 있다고 B 정점에 어떤 정점들이 연결되어 있는지 까지는 알 수 없다는 뜻이다.

그렇다면 A와 B 둘 중 누굴 기준으로 union 해주어야 할까? 바로 친구 비용이 적은 놈을 기준으로 잡아야 한다.

어차피 우리가 원하는 건 준석이가 사용할 최소비용이므로 A의 비용이 B보다 작다면

B->A를 가리키게 하면, 준석이는 A를 가리키든 B를 가리키는 A의 비용만큼 사용할 것이다.

 

union이 끝났으면 친구 관계를 한 명씩 까보는데 인덱스 값이 자기 자신을 가리키고 있는 친구가 있다면

전부 돈을 주고 매수(?)한다. 그렇게 지불한 비용이 초기에 주어진 금액보다 작거나 같으면

사용한 금액을 출력하면 통과할 수 있다.

 

---

3. 코드

import sys

input = sys.stdin.readline
sys.setrecursionlimit(10**8)

def find(num, tree):
    if num == tree[num]:
        return num
    return find(tree[num], tree)

def union(a, b, tree, cost):
    a = find(a, tree)
    b = find(b, tree)

    if a == b:
        return
    else:
        if cost[a] > cost[b]:
            tree[a] = b
        else:
            tree[b] = a

def solution():
    n, m, k = map(int, input().split())
    cost = [0] + list(map(int, input().split()))
    relation = [i for i in range(n+1)]
    for _ in range(m): # friend relation create
        a, b = map(int, input().split())
        union(a, b, relation, cost)

    result = 0
    for idx, root in enumerate(relation):
        if idx == root:
            result += cost[idx]
    
    if result <= k:
        print(result)
    else:
        print("Oh no")

solution()