백준 16562번 <친구비>

백준 알고리즘 16562번

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

16562번: 친구비

---

16562번

친구의 친구는 친구다를 이용하여 n명의 사람들과 모두 친구가 되기 위해서 최소한으로 필요한 친구비를 구하여라.

 

입력으로는 첫 줄에 학생 수 N과 친구 관계 수 M, 가지고 있는 돈 K가 들어온다. 둘째 줄에 N명이 원하는 친구비가 주어진다. 그 다음 M개의 줄에는 숫자 v, w가 주어진다. 이는 v와 w가 서로 친구라는 의미이다.

 

출력으로는 n명과 모두 친구가 되는데 드는 최소 비용을 출력하면 된다.

---

문제 해결

union-find 알고리즘을 이용하여 문제를 해결하였다.

 

find(a)

- a의 root를 찾아준다.

(여기서 root는 해당 친구 그룹에서 친구비가 가장 적은 친구의 값이다.)

 

union(a, b)

- a와 b를 같은 친구 그룹으로 합친다. 이때 친구비가 더 적은 친구의 값으로 update한다.

친구관계가 들어오면 그 둘을 union을 한다. union을 할 때는 해당 친구 그룹에서 친구비가 가장 적은 친구의 값을 갖도록 한다. 친구관계에 대해 모든 작업을 끝내고 나서는 같은 친구 그룹이지만 다른 값을 갖는 것에 대해 친구비가 가장 적은 친구의 값을 갖도록 find를 이용하여 바꿔준다.

필요한 친구비의 합을 구하려면 친구 리스트를 set으로 만들어 같은 친구 그룹인 것을 하나로 만들고 for문을 이용해 해당 친구들의 친구비를 더해서 구한다.

---

코드

import sys
input = sys.stdin.readline
n, m, k = map(int, input().split())
needs = list(map(int, input().split()))
connect = []
for _ in range(m):
    v, w = map(int, input().split())
    connect.append((v-1, w-1))
friends = [i for i in range(n)]
def find(a):
    if friends[a] == a:
        return a
    return find(friends[a])

def union(a, b):
    find_a = find(a)
    find_b = find(b)
    if needs[find_a] < needs[find_b]:
        needs[find_b] = needs[find_a]
    else:
        needs[find_a] = needs[find_b]
    if find_a < find_b:
        friends[find_b] = find_a
    else:
        friends[find_a] = find_b

for v, w in connect:
    union(v, w)

for i in range(n):
    find_i = find(i)
    if find_i != friends[i]:
        friends[i] = find_i
friends = set(friends)
result = 0
for i in friends:
    result += needs[i]
if result > k:
    print("Oh no")
else:
    print(result)

---

결과

---