[Pypy] 백준 20303 할로윈의 양아치

 

 

문제

Trick or Treat!!

10월 31일 할로윈의 밤에는 거리의 여기저기서 아이들이 친구들과 모여 사탕을 받기 위해 돌아다닌다. 올해 할로윈에도 어김없이 많은 아이가 할로윈을 즐겼지만 단 한 사람, 일찍부터 잠에 빠진 스브러스는 할로윈 밤을 즐길 수가 없었다. 뒤늦게 일어나 사탕을 얻기 위해 혼자 돌아다녀 보지만 이미 사탕은 바닥나 하나도 얻을 수 없었다.

단단히 화가 난 스브러스는 거리를 돌아다니며 다른 아이들의 사탕을 빼앗기로 마음을 먹는다. 다른 아이들보다 몸집이 큰 스브러스에게 사탕을 빼앗는 건 어렵지 않다. 또한, 스브러스는 매우 공평한 사람이기 때문에 한 아이의 사탕을 뺏으면 그 아이 친구들의 사탕도 모조리 뺏어버린다. (친구의 친구는 친구다?!)

사탕을 빼앗긴 아이들은 거리에 주저앉아 울고 K$K$명 이상의 아이들이 울기 시작하면 울음소리가 공명하여 온 집의 어른들이 거리로 나온다. 스브러스가 어른들에게 들키지 않고 최대로 뺏을 수 있는 사탕의 양을 구하여라.

스브러스는 혼자 모든 집을 돌아다녔기 때문에 다른 아이들이 받은 사탕의 양을 모두 알고 있다. 또한, 모든 아이는 스브러스를 피해 갈 수 없다.

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

 

 

 

 

 

 

 

난이도

골드3

 

 

 

 

 

 

 

내 코드

import sys

## 유니온 파인드로 그룹 묶기

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


def union(x,y) :
    x = find(x)
    y = find(y)

    if x == y :
        return
    
    if x < y :
        parent[y] = x
    else :
        parent[x] = y

n, m, k = map(int,sys.stdin.readline().split())
candy = [0] + list(map(int,sys.stdin.readline().split()))
parent = [i for i in range(n+1)]
group = []
for _ in range(m) :
    a, b = map(int,sys.stdin.readline().split())
    union(a,b)

for i in range(1,n+1) :
    if parent[i] != i :
        parent[i] = find(i)

weight = [0] * (n+1)
value = [0] * (n+1)

for i in range(1,n+1) :
    p = parent[i]

    weight[p] += 1
    value[p] += candy[i]

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

for i in range(1,n+1) :
    if weight[i] != 0:
        w = weight[i]
        v = value[i]

        if w > k :
            continue
        
        for j in range(k,-1,-1) :
            if j >= w :
                dp[j] = max(dp[j], dp[j-w]+v)

print(dp[k-1])

 

 

 

 

 

풀이

이 문제는 1. 친구 무리 만들기 2. 냅색 알고리즘을 이용해서 답 구하기 단계로 풀었다.

아는 친구들끼리의 그룹을 만들 때 처음에는 BFS를 사용해서 친구들끼리 배열로 만들고, 그 배열의 인원수는 무게로 가진 사탕수는 가치로 두고 2차원 DP를 사용해 냅색알고리즘을 하였으나 시간초과가 났다.

 

수정한 부분은

1. 무리를 만들 때 BFS가 아니라 유니온-파인드 방법을 사용하여 친구인 사람들끼리는 부모가 같도록해서 시간을 단축하였고

2. 2차원 DP 대신 DP를 뒤에서부터 업데이트하는 방법을 이용해 1차원 DP로 최적화하였다.

결과적으로 pypy로는 통과하였으나 python으로는 22% 에서 시간초과가 났다...ㅜㅜ