[Python] 1562번 계단 수

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

1562번: 계단 수

---

 

23/12/05

 

 

비트마스킹을 이용한 DP를 연습하기 좋은 문제이다.

 

클래스 5에도 있으니, 비트마스킹에 대한 개념을 이해한 뒤에 조금 생각해보고 풀면 더욱 도움이 될 것 같다. 

---

 

문제 접근 방식:

 

 

길이가 매우 긴 자리수를 가지고 있는 특정한 계단 수에 대한 경우의 수를 구하는 문제이다.

 

계단 수이긴 하지만, $0$부터 $9$까지의 모든 숫자를 써야 한다는 점이 다르다.

 

즉, 우리는 숫자의 사용 여부를 관리해야 할 필요성이 있다.

 

그걸 DP 배열에 넣고 관리하기에는 조금 까다로운 면이 있다.

 

이럴 때 비트마스킹을 이용할 수 있는데, 사용했는지 하지 않았는지에 대한 여부를 하나의 비트로 관리하여 일종의 "숫자"로 표시하는 것이다.

 

즉, 첫번째 비트는 $0$을 사용했는지의 여부, 두 번째 비트는 $1$을 사용했는지의 여부, ... , 열 번째 비트는 $9$를 사용했는지의 여부로 나타낼 수 있다.

 

사용했다면 그 비트를 1로, 사용하지 않았다면 0으로 놔두면 된다.

 

총 10개의 비트가 사용되므로 우리는 사용되었는지 되지 않았는지에 대한 여부를 1024까지의 수 하나로 관리할 수 있다.

 

 

DP배열은 3차원 배열로 관리할 수 있다.

 

즉, 아래와 같이 정의하면 된다.

 

$$DP[i][j][k] = i \textrm{자리 숫자이고 } j \textrm{로 끝났을 때의 계단 수 개수, } k \textrm{는 각 숫자의 사용 여부}$$

 

이렇게 배열을 정의한다면, 우리는 점화식 자체는 아래와 같이 쉽게 정의할 수 있다.

 

$$DP[i][j][\textrm{after }j] = DP[i-1][j-1][k] + DP[i-1][j][k]$$

 

여기서 $\textrm{after }j$는 $j$번째 비트를 킨 이후를 뜻한다.

 

$k$에서 $j$번째 비트를 키는 연산은 비트 연산을 사용하면 쉽게 구현할 수 있다.

 

비트 쉬프트 연산과 OR연산을 이용하면 $j$번째 비트를 키는 구현은 쉽게 할 수 있다.

 

 

DP배열의 초기 항은 0으로 시작하는 계단 수는 없다고 했으므로, 0을 제외한 한 자리 수(1, ..., 9)와 그에 대한 사용 여부 $k$를 1로 잡아두면 충분하다.

 

즉, $1 \leq i \leq 9$에 대해서 $DP[1][i][1<<i] = 1$로 정의하면 충분하다.

 

DP배열을 갱신할 때에 MOD연산에 주의하며 갱신하는 점을 잊지 말자. 

---

아래는 내가 위의 접근 방식과 같이 작성한 파이썬 코드이다. 더 보기를 누르면 확인할 수 있다.

# 1562번 계단수
# 비트마스킹, DP
'''
접근 방법:
DP[i][j][k] = i자리 숫자이고 j로 끝났을 때 k는 비트마스킹
DP[i][j][k] = DP[i-1][j-1][j번째 비트 키기 이전] + DP[i-1][j][j번째 비트키기 이전]
'''
MOD = 1_000_000_000
DP = [[[0 for _ in range(1024)] for _ in range(10)] for _ in range(101)]
# 0으로 시작하는 숫자는 없다.
for i in range(1, 10):
    DP[1][i][1<<i] = 1
for i in range(2, 101):
    for j in range(10):
        for k in range(1024):
            after = k | (1 << j)
            if j == 0:
                DP[i][j][after] += DP[i-1][j+1][k]
                DP[i][j][after] %= MOD
            elif j == 9:
                DP[i][j][after] += DP[i-1][j-1][k]
                DP[i][j][after] %= MOD
            else:
                DP[i][j][after] += (DP[i-1][j-1][k] + DP[i-1][j+1][k])
                DP[i][j][after] %= MOD
N = int(input())
print(sum(DP[N][i][1023] for i in range(10))%MOD)