[BOJ] 18427 함께 블록 쌓기 (Python3)

문제

1번부터 N번까지의 학생들은 각각 블록들을 가지고 있다. 학생마다 최대 M개의 블록을 가지고 있을 수 있으며, 한 명의 학생이 가지고 있는 모든 블록들의 높이는 서로 다르다. 이 때 1번부터 N번까지의 학생들이 가진 블록을 차례대로 사용하여 바닥에서부터 쌓아올려 하나의 탑을 만들고자 한다.

단, 어떤 학생의 블록은 사용하지 않아도 되며 한 학생당 최대 1개의 블록만을 사용할 수 있다.

1번부터 N번까지의 학생들이 가지고 있는 블록들에 대한 정보가 주어졌을 때, 높이가 정확히 H인 탑을 만들 수 있는 경우의 수를 계산하는 프로그램을 작성하시오.

예를 들어 N=3, M=3, H=5일 때, 각 학생마다 가지고 있는 블록들의 높이가 다음과 같다고 가정하자.

• 1번 학생: 2, 3, 5
• 2번 학생: 3, 5
• 3번 학생: 1, 2, 3

이 때, 탑의 높이가 정확히 5가 되도록 블록을 쌓는 경우로는 다음의 6가지가 존재한다. (블록을 사용하지 않는 경우는 X로 표시하였다.)

1번 학생	2번 학생	3번 학생
x	3	2
x	5	x
2	x	3
2	3	x
3	x	2
5	x	x

입력

첫째 줄에 자연수 N, M, H가 공백을 기준으로 구분되어 주어진다. (1 ≤ N ≤ 50, 1 ≤ M ≤ 10, 1 ≤ H ≤ 1,000) 둘째 줄부터 N개의 줄에 걸쳐서 각 학생이 가진 블록들의 높이가 공백을 기준으로 구분되어 주어진다.

단, 모든 블록의 높이는 1,000 이하의 자연수이며 한 명의 학생이 가지고 있는 모든 블록들의 높이는 서로 다르게 주어진다.

출력

첫째 줄에 높이가 H인 탑을 만드는 경우의 수를 10,007로 나눈 나머지를 출력한다.

시간초과 코드

처음에는 단순하게 dfs로 완전탐색해보려 했다.

하지만 코드를 작성하고 보니 최악의 경우 N * M = 500개의 조합을 구해야하므로 대충 계산해봐도 1억이 훌쩍 넘는다.

따라서 이 코드는 시간초과가 발생했다.

n, m, h = map(int, input().split())

info = dict()
for i in range(n):
    info[i] = list(map(int, input().split()))

for i in info.keys():
    info[i].append(0)
    
answer = 0

def dfs(start, value):
    global answer
    
    if value == h:
        answer += 1
        return
    
    if start == n:
        return
    
 
    if value > h:
        return
    
    for num in info[start]:
        dfs(start + 1, value + num)
    
dfs(0, 0)
print(answer)

그래서 다른 방법을 고민하다가, 이 문제는 결국 개수가 최대 3개 정해져있고 무게가 있는 냅색 문제라는 것을 깨달았다.

그런데 냅색알고리즘을 푼지가 오래되서 기억이 잘 안나 결국 유튜브에서 냅색 알고리즘을 다시 공부했다.

코드없는 프로그래밍님 유튜브 - Knapsack problem

NS(n + 1, h + 1)로 생각하고 이 dp 테이블을 업데이트 해나가면 된다.

행은 학생의 번호를 나타내고, 열은 점수를 나타낸다.

해당 점수를 만들 수 있는 경우의 수를 추가해주며 테이블을 업데이트 해가면 된다.

정답코드

n, m, h = map(int, input().split())

dp = [[0 for _ in range(h + 1)] for _ in range(n + 1)]

for i in range(n + 1):
    dp[i][0] = 1
    
for i in range(1, n + 1):
    info = list(map(int, input().split()))
    
    dp[i] = dp[i-1].copy()
    
    for j in info:
        for k in range(j, h + 1):
            dp[i][k] += dp[i - 1][k - j]
            
print(dp[n][h] % 10007)