Baekjoon 20303. 할로윈의 양아치 / Python

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

알고리즘 설명

문제 설명

스브러스는 할로윈 밤에 혼자서 돌아다니며 다른 아이들의 사탕을 빼앗기로 결정했다.
스브러스는 아이들의 친구 관계를 알고 있으며, 한 아이의 사탕을 빼앗으면 그 아이의 친구들의 사탕도 모두 빼앗는다.
스브러스가 최대로 뺏을 수 있는 사탕의 양을 구하는 것이 목표다.
단, K명 이상의 아이들이 울면 어른들이 나와서 스브러스를 멈추게 합니다.

코드 설명

1. Union-Find: 친구 관계를 통해 연결된 모든 아이들을 하나의 그룹으로 묶는다.
2. 그룹 내 사탕과 친구 수 집계: 각 그룹별로 총 사탕의 수와 친구의 수를 계산한다.
3. DP (동적 계획법): 각 그룹을 대상으로 하여 그룹을 선택하거나 선택하지 않는 경우를 고려하여 최대 사탕 수를 계산한다. K 이상의 아이들이 울지 않도록 K명을 초과하지 않게 조절한다.

시간 복잡도

1. Union-Find: O(α(N)) = 거의 상수 시간
2. 동적 계획법: O(N×K)

import sys
input = sys.stdin.readline

# Union-Find 함수들
def union_find(x):
    if parent[x] != x:
        parent[x] = union_find(parent[x]) # 경로 압축
    return parent[x]

def union_set(x, y):
    rootX = union_find(x)
    rootY = union_find(y)
    if rootX != rootY:
        parent[rootY] = rootX # union 연산

# DP 함수
def rob(s):
    for j in range(k - 1, friends_num[s] - 1, -1):
        dp[j] = max(dp[j], dp[j - friends_num[s]] + candies[s])

n, m, k = map(int, input().split()) # 아이들의 수, 친구 관계 수, 울음소리가 공명하기 위한 최소 아이의 수
candies = [0] + list(map(int, input().split())) # 각 아이가 받은 사탕의 수
parent = [i for i in range(n + 1)] # Union-Find의 parent 배열

# 친구 관계 입력받아 Union-Find로 묶기
for _ in range(m):
    a, b = map(int, input().split())
    union_set(a, b)

# 그룹별 사탕 수와 친구 수 계산
friends_num = [1] * (n + 1)
dp = [0] * (k)

for i in range(1, n + 1):
    root = union_find(i)
    if root != i:
        candies[root] += candies[i]
        friends_num[root] += friends_num[i]

# 그룹별로 DP 수행
for i in range(1, n + 1):
    if i == parent[i]: # 루트 노드인 경우만
        rob(i)

# 결과 출력
print(max(dp))