[C++] 계단 수 (백준 1562번)

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

 

10844번 쉬운 계단 수 문제의 어려운 버전인가 보다. 그 문제에 대한 풀이는 이 글에 있다. 

쉬운 계단 수와의 다른 점은 0부터 9까지의 수가 모두 등장해야 한다는 것이다. 그러면 어떤 어떤 수를 사용했는지 비트마스킹으로 표시해주면 되겠네!

 

dp[i][j][k]: 길이가 i이면서 j로 끝나는데, 현재까지 사용한 수가 k 상태인 계단 수의 개수

(0을 사용했으면 2^0으로, 1을 사용했으면 2^1으로... 표현하고 그걸 모두 합한 값을 k로 사용한다. 즉 k는 0 이상 1023 이하의 값이다.)

0으로 끝나는 수 → dp[i][0][k | 2^0] += dp[i - 1][1][k]

9로 끝나는 수 → dp[i][j] = dp[i][9][k | 2^9] += dp[i - 1][8][k]

나머지로 끝나는 수 → dp[i][j][k | 2^j] += (dp[i - 1][j - 1][k] + dp[i - 1][j + 1][k])

위의 점화식을 사용하되, mod 연산만 중간중간에 잘 해주면 된다.

 

#include <iostream>
#include <vector>
using namespace std;

int mod = 1000000000;
int pow2[] = {1, 2, 4, 8, 16, 32, 64, 128, 256, 512};

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    int n;
    cin >> n;

    vector<vector<vector<int>>> dp(n + 1, vector<vector<int>>(10, vector<int>(1024, 0)));

    for (int j = 1; j < 10; j++) {
        dp[1][j][pow2[j]] = 1;
    }
    
    for (int i = 2; i < n + 1; i++) {
        for (int k = 0; k < 1024; k++) {
            dp[i][0][k | pow2[0]] = (dp[i][0][k | pow2[0]] + dp[i - 1][1][k]) % mod;
            for (int j = 1; j < 9; j++) {
                dp[i][j][k | pow2[j]] = (dp[i][j][k | pow2[j]] + dp[i - 1][j - 1][k]) % mod;
                dp[i][j][k | pow2[j]] = (dp[i][j][k | pow2[j]] + dp[i - 1][j + 1][k]) % mod;

            }
            dp[i][9][k | pow2[9]] = (dp[i][9][k | pow2[9]] + dp[i - 1][8][k]) % mod;
        }
    }

    int ans = 0;
    for (int j = 0; j < 10; j++) ans = (ans + dp[n][j][1023]) % mod;

    cout << ans << '\n';
}

(AC)