백준풀이

[Python] 백준 16562 친구비

SoU330 2024. 10. 2. 04:49

 

문제

19학번 이준석은 학생이 N명인 학교에 입학을 했다. 준석이는 입학을 맞아 모든 학생과 친구가 되고 싶어한다. 하지만 준석이는 평생 컴퓨터랑만 대화를 하며 살아왔기 때문에 사람과 말을 하는 법을 모른다. 그런 준석이에게도 희망이 있다. 바로 친구비다!

학생 i에게 Ai만큼의 돈을 주면 그 학생은 1달간 친구가 되어준다! 준석이에게는 총 k원의 돈이 있고 그 돈을 이용해서 친구를 사귀기로 했다. 막상 친구를 사귀다 보면 돈이 부족해질 것 같다는 생각을 하게 되었다. 그래서 준석이는 “친구의 친구는 친구다”를 이용하기로 했다.

준석이는 이제 모든 친구에게 돈을 주지 않아도 된다!

위와 같은 논리를 사용했을 때, 가장 적은 비용으로 모든 사람과 친구가 되는 방법을 구하라.

 

 

 

 

 

 

난이도

골드 4

 

 

 

 

 

 

내 코드

import sys

def find(x) :
    if parent[x] != x :
        parent[x] = find(parent[x])
        return parent[x]
    return x

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

    if a < b :
        parent[b] = a
    else :
        parent[a] = b

N, M, k = map(int,sys.stdin.readline().split())
cost = [0] + list(map(int,sys.stdin.readline().split()))
parent = [i for i in range(N+1)]
pay = dict()
answer = 0

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

    v = find(v)
    w = find(w)
    union(v,w)

for i in range(1,N+1) :
    group = find(i)
    friend_cost = cost[i]

    if group in pay :
        if pay[group] > friend_cost :
            answer -= pay[group]
            answer += friend_cost
            pay[group] = friend_cost
    else :
        pay[group] = friend_cost
        answer += friend_cost

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

 

 

 

 

 

 

풀이과정

유니온-파인드를 사용할 줄 알면 쉽게 풀리는 문제다.

부모를 찾고 최적화하는 find() 함수와 두 개를 합치는 union() 함수를 구현해둔다.

여기서 친구관계를 입력받을 때 마다 union() 함수를 이용해 부모를 갖게한다.

최종적으로 부모가 같은 친구들은 모두 같은 무리이기 때문에 그 무리들 중 가장 친구비가 적은 친구와 친해지면 된다.

최소비용을 구할 때는 딕셔너리를 사용해서 딕셔너리에 부모 번호와 현재까지 나온 수 중 제일 작은 수를 넣어두어 처리하였다.