BOJ 20303 파이썬

문제

Trick or Treat!!
10월 31일 할로윈의 밤에는 거리의 여기저기서 아이들이 친구들과 모여 사탕을 받기 위해 돌아다닌다. 올해 할로윈에도 어김없이 많은 아이가 할로윈을 즐겼지만 단 한 사람, 일찍부터 잠에 빠진 스브러스는 할로윈 밤을 즐길 수가 없었다. 뒤늦게 일어나 사탕을 얻기 위해 혼자 돌아다녀 보지만 이미 사탕은 바닥나 하나도 얻을 수 없었다.
단단히 화가 난 스브러스는 거리를 돌아다니며 다른 아이들의 사탕을 빼앗기로 마음을 먹는다. 다른 아이들보다 몸집이 큰 스브러스에게 사탕을 빼앗는 건 어렵지 않다. 또한, 스브러스는 매우 공평한 사람이기 때문에 한 아이의 사탕을 뺏으면 그 아이 친구들의 사탕도 모조리 뺏어버린다. (친구의 친구는 친구다?!)
사탕을 빼앗긴 아이들은 거리에 주저앉아 울고
$K$명 이상의 아이들이 울기 시작하면 울음소리가 공명하여 온 집의 어른들이 거리로 나온다. 스브러스가 어른들에게 들키지 않고 최대로 뺏을 수 있는 사탕의 양을 구하여라.
스브러스는 혼자 모든 집을 돌아다녔기 때문에 다른 아이들이 받은 사탕의 양을 모두 알고 있다. 또한, 모든 아이는 스브러스를 피해 갈 수 없다.

입력
첫째 줄에 정수 $N$, $M$, $K$가 주어진다. $N$은 거리에 있는 아이들의 수, $M$은 아이들의 친구 관계 수, $K$는 울음소리가 공명하기 위한 최소 아이의 수이다. ($1 \leq N \leq 30\ 000$, $0 \leq M \leq 100\ 000$, $1 \leq K \leq \min\left\{N, 3\ 000\right\}$)

둘째 줄에는 아이들이 받은 사탕의 수를 나타내는 정수 $c_1, c_2, \cdots, c_N$이 주어진다. ($1 \leq c_i \leq 10\ 000$)

셋째 줄부터
$M$개 줄에 갈쳐 각각의 줄에 정수 $a$, $b$가 주어진다. 이는 $a$와 $b$가 친구임을 의미한다. 같은 친구 관계가 두 번 주어지는 경우는 없다. ($1 \leq a, b \leq N$, $a \neq b$)

출력
스브러스가 어른들에게 들키지 않고 아이들로부터 뺏을 수 있는 최대 사탕의 수를 출력한다.

접근

1. 친구관계가 서로소집합으로 나누어지고 각각의 집단에 대해서 인원수가 weight, 각 인원이 갖고 있는 사탕 수의 합이 value가 되는 냅색문제라고 판단
2. 처음에 DFS를 이용하여 친구관계를 그룹핑함. (시간초과) 유니온 파인드 알고리즘을 사용해야 한다는 것을 알게 됨.

코드

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**7)

# 유니온 파인드
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[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, input().split())
k -= 1
values = [0] + list(map(int, input().split()))

parent = [x for x in range(n+1)]
for _ in range(m):
    a, b = map(int, input().split())
    union(a, b)

groups = [0] * (n+1)
for i in range(1, n+1):
    if i != parent[i]:
        I = find(i)
        values[I] += values[i]
        groups[I] += 1
    else:
        groups[i] += 1


# 냅색 알고리즘
table = [0] * (k+1)

for i in range(1, n+1):
    if groups[i] != 0:
        w = groups[i]
        v = values[i]
        if w > k:
            continue
        for j in range(k, 0, -1):
            if j + w <= k and table[j] != 0:
                table[j+w] = max(table[j+w], table[j] + v)
        table[w] = max(table[w], v)

print(max(table))

union find 알고리즘

유니온 파인드 알고리즘은 크게 두 개의 함수로 이루어져 있다.
1. find 함수

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

x 노드의 부모 노드가 무엇인지 find해주는 함수이자, 트리의 높이를 낮게 유지해주는 함수이다.

2. union 함수

def union(a, b):
    a = find(a)
    b = find(b)
    if a<b:
        parent[b] = a
    else:
        parent[a] = b

union 함수는 a, b 부모 중 더 작은 쪽으로 합쳐주어 합집합 연산을 해주는 함수이다.

냅색 알고리즘

# 냅색 알고리즘
table = [0] * (k+1)

for i in range(1, n+1):
    if groups[i] != 0:
        w = groups[i]
        v = values[i]
        if w > k:
            continue
        for j in range(k, 0, -1):
            if j + w <= k and table[j] != 0:
                table[j+w] = max(table[j+w], table[j] + v)
        table[w] = max(table[w], v)

냅색 알고리즘은 최대 k 무게를 넘지 않도록 각각 무게와 가치가 다른 보석을 담는 문제로 유명하다.
유니온 파인드 알고리즘으로 그룹을 만들고, 여기서 무게는 각 그룹의 인원 수이고 가치는 그룹에 포함된 학생들이 가지고 있는 사탕 수의 합이다.

냅색 알고리즘은 1~k까지 무게별로 넣을 수 있는 최대 가치를 기록해두는 dp를 활용한다.
보석을 하나씩 비교하는 방식이다. 매번 해당 보석의 무게를 더했을 때 해당되는 dp에 최대 가치를 저장한다.