백준 - 16562번: 친구비 (Python, 분리 집합)

목차

• 문제명
• 문제풀이
• 코드

 

친구비 

 

문제 유형 : 그래프 이론, 그래프 탐색, 분리 집합

난이도 : 골드 4

 

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

 

 

 

문제 풀이

 

'친구의 친구는 친구' 이기 때문에 분리 집합으로 친구끼리 같은 집합으로 묶은 뒤, 각 집합에서 최소 Ai만 뽑아내어 합했을때 값이 K 이하면 출력, K 초과면 "Oh no"를 출력하면 된다.

 

구현 시 주의사항은, union 메소드로 요소가 다수인 집합과 요소가 다수인 집합을 묶었을때 집합의 모든 root값을 업데이트 해주어야 한다.

 

예로,

 

root[1] = 1

 

root[2] = 2

root[3] = 2

root[4] = 2

 

root[5] = 1

 

이고 union 함수가 

 

def union(a, b):
    a = find_root(a)
    b = find_root(b)
    if a > b:
        root[a] = b
    elif a < b:
        root[b] = a

 

위와 같은 경우 

 

4, 5가 친구라는 입력이 주어지면 

 

root[1] = 1

 

root[2] = 1

root[3] = 2

root[4] = 2

 

root[5] = 1

 

이 되므로 root를 업데이트 해주는 과정이 필요하다.

 

코드

 

import sys
from collections import deque
input = sys.stdin.readline

N, M, k = map(int, input().split())
Ai = deque(map(int, input().split()))
root = [i for i in range(N + 1)]
Ai.appendleft('empty')


def find_root(a):
    if root[a] == a:
        return a
    root[a] = find_root(root[a])
    return root[a]


def union(a, b):
    a = find_root(a)
    b = find_root(b)
    if a > b:
        root[a] = b
    elif a < b:
        root[b] = a

# 친구 사이는 같은 집합으로 묶기
for _ in range(M):
    a, b = map(int, input().split())
    union(a, b)

# 모든 root값 업데이트 과정
for i in range(1, N + 1):
    root[i] = find_root(root[i])

queue = []
# (조상, 친구비)로 배열에 저장
for i in range(1, N + 1):
    r = root[i]
    cost = Ai[i]
    queue.append((r, cost))
    queue.sort(key=lambda x: (x[0], x[1])) # 조상, 친구비 기준으로 오름차순 정렬

prev_root = queue[0][0]
total = queue[0][1]

for i in range(1, N):
    r, cost = queue[i]
    if prev_root == r: # 같은 조상이면 넘어가기
        continue
    else:	# 조상이 다르다면 친구비 합산하고 조상 바꾸기
        prev_root = r
        total += cost

    if total > k: # 친구비가 k 초과하면 Oh no
        print("Oh no")
        sys.exit(0)

print(total)

 

 

마무리