[백준 16562] 친구비 - Python

문제 보기

 

[사용한 알고리즘]

Union-Find

 

[알고리즘]

1. 친구의 친구는 친구이기 때문에 서로 친구들인 모임을 Union-Find를 이용해서 찾습니다.

2. 해당 모임에서 친구 비용이 가장 작은 학생을 친구로 만듭니다.

3. 위 과정을 모든 학생과 친구를 맺을 때까지 반복합니다.

4. 모든 학생과 친구를 맺는 비용이 현재 가지고 있는 비용보다 적으면 친구로 만드는데 드는 최소비용을 출력합니다. 만약 현재 가지고 있는 비용이 모든 학생과 친구를 맺는 비용보다 적다면 Oh no를 출력합니다.

 

[코드]

import sys


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


def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b


if __name__ == "__main__":
    n, m, k = map(int, sys.stdin.readline().split())
    friends = list(map(int, sys.stdin.readline().split()))
    parent = [i for i in range(n + 1)]

    # 친구 관계 구현
    for _ in range(m):
        v, w = map(int, sys.stdin.readline().split())
        union_parent(parent, v, w)

    # 친구들의 모임 파악 후 친구 형성
    answer = 0
    visited = set()
    for i in range(1, n + 1):
        if i not in visited:  # 아직 친구가 아닌 모임
            temp = [i]
            cost = friends[i - 1]
            for j in range(1, n + 1):
                if i != j:
                    if find_parent(parent, i) == find_parent(parent, j):
                        temp.append(j)
                        cost = min(cost, friends[j - 1])
            visited.update(temp)
            answer += cost

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