[Python] 백준 1562번: 계단 수

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

1562번: 계단 수

 

import sys
input = sys.stdin.readline
MOD = 1000000000
MAX = (1 << 10)-1


def get_stair_count():
    dp = [[0]*(MAX+1) for _ in range(10)]

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

    for _ in range(2, N+1):
        next_dp = [[0]*(MAX+1) for _ in range(10)]

        for i in range(10):
            for j in range(MAX+1):
                if i > 0:
                    next_dp[i][j | (1 << i)] = (
                        next_dp[i][j | (1 << i)] + dp[i-1][j]) % MOD
                if i < 9:
                    next_dp[i][j | (1 << i)] = (
                        next_dp[i][j | (1 << i)] + dp[i+1][j]) % MOD
        dp = next_dp

    return sum(dp[i][MAX] for i in range(10)) % MOD


if __name__ == '__main__':
    N = int(input())
    print(get_stair_count())

 

결론부터 말하면, DP[길이][마지막 자리의 숫자][BITMASKING]를 통해 문제를 풀었다.

다만 문제를 풀어보면 알겠지만, DP는 바로 이전의 것만을 참조한다. 따라서 DP[마지막 자리의 숫자][BITMASKING]을 사용하고, 이것을 2개 저장하면 된다.

 

마지막 자리의 숫자를 저장하는 이유는, 계단 수를 위해서이다. get_stair_count() 함수는, 마지막 자리의 수에 마지막 자리의 수와 인접한 수를 추가하는 방식으로 진행되기 때문이다. 그러니까, 123456789에서 9가 마지막 자리의 숫자이고, 다음에는 여기에 8을 더하는 식으로 간다는 것이다.

BITMASKING은 각 자리에 1번이라도 숫자가 들어갔는지를 체킹한다. 예를 들어, 345에 대한 비트마스킹은, 0001110000이다.

 

이렇게 한다는 발상을 하는 것이 매우 어려웠다. 마지막 자리의 숫자를 넣는다는 것은, 계단수를 끝에 1개씩 넣어주면서 진행해야겠다는 생각을 기반으로 둔 것이다. 그리고 비트마스킹으로 각 숫자가 있는지를 표시하는 것은, 문제의 조건과 매우 밀접해있었다.

 

이것을 기반으로 get_stair_count()을 설명한다.

먼저 길이가 1에 해당하는 계단수에 대한 DP를 초기화시킨다. 다만, 0으로 시작할 수 없으므로 1~9까지 진행한다.

다음으로 삼중반복문을 수행한다. 2~N까지 for문을 진행하는데, 이것은 현재 길이라고 생각하면 쉽다.(사용하지는 않는다.)

내부를 살펴보면, 마지막 자리의 숫자(i)에 따라 next_dp에 값을 추가한다. 마지막 자리의 숫자에서 숫자를 1개씩 더하는 식으로 진행이 되므로, 이렇게 조건이 붙는 것이다. 가령 마지막 자리 숫자가 9이면 다음에는 8이 올 수밖에 없지만, 6이라면 다음에 5, 7이 와야한다. 아무튼, 이전에 저장되었던 dp를 참조하여 next_dp에 값을 넣어준다. dp, next_dp를 잘 구분해주면 어려울 것이 없다고 생각한다.

 

마지막으로 우리가 원하는 값을 dp를 통해 리턴한다. 좀 쉽게 풀어서 쓰자면, dp[0~9][1111111111]에 해당하는 것들의 합을 리턴해주면 된다.