[Python] 계단 수 - 백준 (DP/비트마스킹)

해당 글은 백준 1562번 '계단 수' 문제를 다룬다.

문제

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

1562번: 계단 수

 

설명

해당 문제를 처음 봤을 때, 예전에 풀어본 문제와 비슷한거 같아 찾아보니 이전의 쉬운 계단 수 문제를 풀었었다.

쉬운 계단 수 문제와 다른 점은 0~9까지 숫자가 어떤 수가 포함되어 있는지를 확인해야 한다.

이때 비트마스킹을 이용한 DP로 이진수로 표현해서 연산하는 방법이 있다.

먼저 0~9까지를 다 포함하는 계단 수를 구하려면 dp 배열이 (n+1) x 10 x 1023 이 되어야 한다.

다음 한자리 수는 다 1로 초기화하고, 다음 자리수부터 3중 for문으로 구한다.

쉬운 계단 수에서 0~1023까지의 총 1024가지를 비트로 표현해서 DP를 진행하면 된다.

 

더 자세하게 잘 설명해준 참고링크를 남긴다.

*참고링크

[백준(BOJ)] #1562- 계단 수 (파이썬, PyPy3)

 

코드

n = int(input())
bit_range = 1 << 10    # 1024 - 0~9까지 사용할 수 있는 숫자 리스트를 포함하기 위함

dp = [[[0]*(bit_range) for _ in range(10)] for _ in range(n+1)]

for i in range(1, 10):
    dp[1][i][1 << i] = 1

for i in range(2, n+1):
    for j in range(10):
        for k in range(bit_range):
            temp = k | 1 << j
            if j == 0:
                dp[i][j][temp] += dp[i-1][j+1][k]
            elif j == 9:
                dp[i][j][temp] += dp[i-1][j-1][k]
            else:
                dp[i][j][temp] += dp[i-1][j-1][k] + dp[i-1][j+1][k]
result = 0
for i in range(10):
    result += dp[n][i][bit_range-1]
print(result%1000000000)