Always awake

[백준][20303][Python 파이썬][G3] 할로윈의 양아치 본문

PS/백준 BOJ

[백준][20303][Python 파이썬][G3] 할로윈의 양아치

죠. 2023. 12. 8. 19:41

 

문제

Trick or Treat!!

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

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

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

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

입력

첫째 줄에 정수 N, M, K가 주어진다.N은 거리에 있는 아이들의 수, M은 아이들의 친구 관계 수, K는 울음소리가 공명하기 위한 최소 아이의 수이다.

둘째 줄에는 아이들이 받은 사탕의 수를 나타내는 정수가 주어진다. 

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

 

출력

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

 

테스트케이스

10 6 6
9 15 4 4 1 5 19 14 20 5
1 3
2 5
4 9
6 2
7 8
6 10
5 4 4
9 9 9 9 9
1 2
2 3
3 4
4 5
10 9 11
1 2 3 4 5 6 7 8 9 10
1 2
2 3
1 4
5 7
6 8
4 5
9 10
8 9
7 10
5 2 2
1 2 4 8 16
1 4
2 3
57 0 55 16

 

코드

import sys
input = sys.stdin.readline

N, M, K = map(int, input().split())  # 아이들 수, 친구 관계 수, 공명 시작 최소 아이수

candylst = [0] + [*map(int, input().split())]  # 사탕 개수들
childs = [*range(N+1)]  # 일단 아이들의 친구는 자기 자신만
friends = [0] + [1 for i in range(N)]  # 각 아이에 대한 친구 집합의 크기


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


def merge(a, b):
    aroot = root(a)
    broot = root(b)
    if aroot != broot:
        if aroot > broot:  # 숫자가 작은쪽으로 몰자
            friends[broot] += friends[aroot]  # 친구 개수 늘려주기
            candylst[broot] += candylst[aroot]
            childs[aroot] = broot
        else:
            friends[aroot] += friends[broot]
            candylst[aroot] += candylst[broot]
            childs[broot] = aroot
        return False  # 둘은 친구관계가 아니었다
    return True  # 둘은 친구관계였다


for i in range(M):
    a, b = map(int, input().split())
    merge(a, b)

selections = []
for i in range(1, N+1):
    if i == childs[i]:
        selections.append([friends[i], candylst[i]])

# print(selections)

bag = [[0 for i in range(len(selections) + 1)]
       for j in range(K)]  # 0~K-1로 K개, K는안되니까

for i in range(len(selections)):
    friendcnt, candycnt = selections[i]
    for j in range(1, K):
        if friendcnt <= j:
            bag[j][i+1] = max(bag[j][i], bag[j-friendcnt][i] + candycnt)
        else:
            bag[j][i+1] = max(bag[j-1][i+1], bag[j][i])


# for i in range(K):
#     print(bag[i])

print(bag[K-1][len(selections)])

 

후기

분리집합을 쓰면서..
각 집합이 몇명인지 체크하면서
그 집합 안에있는 사탕 개수의 총합을 체크한다음에
배낭문제로 풀면 되겠다!

 

라고 생각하고 들어갔고, 배낭문제 오랜만에 쓰니까 조금 어려웠다.

시간이 아슬아슬하던데, 다른 사람들 보니까 배낭문제를 1차원으로 풀던데

selections = []
for i in range(1, N+1):
    if i == childs[i]:
        selections.append([friends[i], candylst[i]])

selections.sort(key=lambda x: (x[0], x[1]))

bag = [0] * (K+1)
for friendcnt, candycnt in selections:
    for i in range(K, friendcnt-1, -1):
        bag[i] = max(bag[i-friendcnt]+candycnt, bag[i])

print(bag[K-1])

 

이렇게 배낭문제 하는 방법을 익혔다. 어차피 해당 배낭에 대해서 그 위치에서 최대만 기록하면 되니까 이렇게 되는구나

이것만으로 시간이 1/10으로.. (3900ms -> 390ms) 꼭 익혀야겠다

 

 

물론 우리 느린 파이썬친구는 그래도 시간초과다