[백준 1415번] 사탕 (Python3)

import sys,math,collections
input = sys.stdin.readline

N = int(input())
candy = [int(input()) for i in range(N)]
S = sum(candy)
candy = collections.Counter(candy); N = len(candy)

DP = [[0]*(S+1) for i in range(N+1)]; DP[0][0] =  1

n = 0
for c in candy:
  for s in range(S+1):
    DP[n+1][s] = DP[n][s]
    for i in range(1,candy[c]+1):
      if s-i*c<0:
        break
      DP[n+1][s] += DP[n][s-i*c]
  n += 1

prime = []; check = [0]*(S+1); sq = int(math.sqrt(S))
for p in range(2,S+1):
  if check[p]:
    continue
  prime.append(p)
  if p>sq:
    continue
  i = 2
  while i*p<=S:
    check[i*p] = 1
    i += 1

result = 0
for p in prime:
  result += DP[N][p]
print(result)

냅색 알고리즘으로 해결하였다.

 

풀이가 잘 떠오르지 않는 문제였다. 골드 문제인데도 불구하고 꽤 오랫동안 풀지 못했던 문제였는데, 이 문제가 냅색 문제와 관련되었다는걸 알게되자 쉽게 해결할 수 있었다.

 

풀이를 설명하면,

1. 각 사탕의 개수를 파이썬의 collections.Counter 함수를 사용해서 세어준다. (중복방지)

2. 냅색 알고리즘으로 각 사탕을 넣었을 때 방법이 존재하는 가짓수를 세어준다.

3. 에라토스테네스의 체 방식으로 소수를 구한다.

4. 2 에서 찾은 가짓수 중 소수에 해당되는 가짓수를 출력한다.

 

그동안 냅색 문제는 문제에서 대놓고 "나 냅색문제요" 하는 문제를 제외하면 풀어본 경험이 없었는데, 이렇게 냅색임을 숨기고 응용해야하는 문제가 있을 수 있다는 것을 알게되었다.

 

 

오늘의 교훈) 냅색 알고리즘을 풀이에 활용하자