[Python/파이썬] 백준 15789번 CTP 왕국은 한솔 왕국을 이길 수 있을까?

15789번: CTP 왕국은 한솔 왕국을 이길 수 있을까?

 

문제

CTP 왕국은 정말 깊은 역사를 가지고 있다. 선대 김진서 왕부터 시작하여 전현용 왕을 거쳐 … 마침내 김세진이 CTP 왕국의 왕이되었다. 세진이는 재미없는 개그를 정말 싫어했기 때문에 왕이 되자마자 CTP 왕국에서 가장 재미없는 이한솔을 쫓아냈다.

화가난 한솔이는 자기의 개그에 유일하게 웃어주던 박정률과 함께 한솔 왕국을 세웠다.

그 이후 33년이 지났다 …………. 

어느새 한솔 왕국은 번창하여 CTP 왕국보다 힘이 쎄졌다. 세진이는 다른 왕국과 동맹을 맺어 CTP 왕국의 힘을 길러 한솔 왕국보다 부흥시키려고 한다.  왕국의 힘이란 동맹국의 수를 의미한다.  (예를 들어 동맹이 없는 나라의 힘은 1이다)

왕국간의 동맹의 법칙은 조금 특별해서 만약에 A왕국과 B왕국이 동맹이고 B왕국과 C왕국이 동맹이라면 A왕국과 C왕국도 동맹이 된다. 

CTP 왕국의 왕 세진이는 최대 K번 다른 왕국과 동맹을 맺을 기회를 갖으며, 현재 동맹관계는 CTP 왕국과 한솔 왕국은 동맹이 아니다. 또한 한솔 왕국과 동맹인 왕국과는 동맹을 맺을 수 없으며 K번의 동맹 맺을 기회를 모두 사용하지 않아도 된다.

각 왕국들의 동맹관계와 CTP 왕국의 번호, 한솔 왕국의 번호가 주어질 때 세진이를 도와 CTP 왕국의 힘의 최댓값을 구하여라. 각 왕국의 번호는 1부터 N까지의 자연수로 나타내어지며, 서로 다른 두 왕국이 같은 번호를 갖는 경우는 없다.

입력

입력의 첫째 줄에 왕국의 수 N(3 ≤ N ≤ 100,000)과 동맹 관계의 수 M(1 ≤ M ≤ 200,000)이 주어진다. 이 후 M개의 줄에 X,Y가 주어진다. 이는 X 왕국과 Y 왕국이 동맹이라는 뜻이다.

입력의 마지막 줄에 CTP 왕국의 번호 C와 한솔 왕국의 번호 H와 추가 동맹의 기회 K(0 ≤ K ≤ 100)가 공백으로 구분되어 주어진다. 

주어지는 입력에서 CTP 왕국과 한솔 왕국은 절대로 동맹이 되지 않게 주어진다.

출력

CTP 왕국의 힘의 최댓값을 출력하라.

코드

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

n, m = map(int, input().split())
graph = [[] for _ in range(n+1)]
for _ in range(m):
    x, y = map(int, input().split())
    graph[x].append(y)
    graph[y].append(x)
c, h, k = map(int, input().split())


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


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


parent = [i for i in range(n+1)]
for a in range(1, n+1):
    for b in graph[a]:
        if find_parent(a, parent) == find_parent(b, parent):
            continue
        else:
            union_parent(a, b, parent)

group = defaultdict(int)    # 동맹
ctp_root, hs_root = -1, -1
for i in range(1, n+1):
    root = find_parent(i, parent)
    group[root] += 1
    if i == c:
        ctp_root = i
    if i == h:
        hs_root = i

ans = group[ctp_root]
for key, val in sorted(group.items(), key=lambda x: x[1], reverse=True):
    if key == hs_root or key == ctp_root:
        continue
    if k <= 0:  # 동맹 맺을 기회 모두 사용
        break
    ans += val
    k -= 1
print(ans)

 

parent 배열은 각 왕국이 속한 그룹의 루트가 되는 왕국의 번호를 저장하는 배열이다.

group 딕셔너리의 key는 그룹의 루트가 되는 왕국의 번호이고 value는 그룹에 속한 왕국들의 수이다. CTP 왕국은 동맹을 최대 k번 맺어서 왕국의 힘을 최대로 만들어야 하므로 소속된 왕국의 수가 큰 그룹부터 동맹을 맺어야 한다. 이때 이미 CTP왕국과 동맹을 맺었거나 한솔 왕국과 동맹인 그룹과는 동맹을 맺을 수 없다.