P.1562 계단 수

코드

import java.io.*;

public class Main {

    static final int MOD = 1_000_000_000;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        int n = Integer.parseInt(br.readLine());
        int[][][] dp = new int[n + 1][10][1024];
        for (int i = 1; i <= 9; i++) dp[1][i][1 << i] = 1;
        for (int i = 2; i <= n; i++) {
            for (int j = 0; j <= 9; j++) {
                for (int k = 0; k < 1024; k++) {
                    int bit = k | (1 << j);
                    if (j == 0) dp[i][j][bit] = (dp[i][j][bit] + dp[i - 1][j + 1][k]) % MOD;
                    else if (j == 9) dp[i][j][bit] = (dp[i][j][bit] + dp[i - 1][j - 1][k]) % MOD;
                    else dp[i][j][bit] = (dp[i][j][bit] + dp[i - 1][j - 1][k] + dp[i - 1][j + 1][k]) % MOD;
                }
            }
        }

        long ans = 0L;
        for (int i = 0; i <= 9; i++) {
            ans = (ans + dp[n][i][1023]) % MOD;
        }
        bw.write(Long.toString(ans));
        bw.flush();
    }
}

코드 설명

문제에서 원하는 것을 구하기 위해서는 다음 네 가지 정보가 필요하다.

1. 자릿수
2. 최소 수
3. 최대 수
4. 마지막 수

그리고 하나의 계단 수를 만들 때 마지막 수를 추가하는 과정에서 오름차순, 내림차순 두 가지 경우가 있다고 판단했다.

(이전 수가 1일 때, 10 혹은 12가 만들어질 수 있음)

 

위 조건들을 모두 포함한 초기 점화식은 다음과 같았다.

dp[a][b][c][d] = dp[a - 1][b + 1][c][d + 1] + dp[a - 1][b][c - 1][d - 1]

dp[a][b][c][d] = 자릿수가 a, 최소 수가 b, 최대 수가 c, 마지막 수가 d인 계단 수의 개수

dp[a - 1][b + 1][c][d + 1] = 내림차순 (이전 최소 수가 b + 1이며 마지막 수가 d + 1)

dp[a - 1][b][c - 1][d - 1] = 오름차순 (이전 최대 수가 c - 1이며 마지막 수가 d - 1)

 

하지만 배열의 요소가 많아지면서 범위 체크할 부분이 너무 많아 알고리즘을 변경했다.

위 점화식은 검증이 되지 않았기 때문에 개념만 가져왔다고 보면 좋을 것 같다.

 

다음 아이디어는 위 점화식에서 최소 수와 최대 수를 합쳐 비트로 보았다.

비트로 볼 경우, 최소 수나 최대 수 뿐만 아니라 현재 dp에서 어떤 수들을 가지고 있는지 확인할 수 있다.

 

dp[a][b][c] = a 자리이며 마지막 수가 b인 계단 수가 가지고 있는 숫자 c(비트)

일단, dp의 크기는 dp[100][10][1024]이다.

c는 비트를 가리키는데 _ _ _ _ _ _ _ _ _ _ 자리가 있고 각각의 자리는 숫자 하나를 나타낸다.

예를 들어 0이 존재한다면 c는 _ _ _ _ _ _ _ _ _ 1이 되며, 이 상태에서 6이 존재한다면 c는 _ _ _ 1 _ _ _ _ _ 1이 된다.

이 비트는 각 자릿수마다 1024를 모두 돌면서 마지막 자릿수를 추가한다.

초기 점화식은 필요한 것만 확인해서 보는 느낌이면 비트마스킹을 이용한 점화식은 계단 수가 만들어질 수 없는 조건이어도 비트를 모두 확인해보는 느낌이다.

for (int i = 2; i <= n; i++) {
    for (int j = 0; j <= 9; j++) {
        for (int k = 0; k < 1024; k++) {
            int bit = k | (1 << j);
        }
    }
}

 

초기 점화식에서 가져온 개념을 대입해보면 다음과 같다.

dp[i][j][bit] = dp[i][j][bit] + dp[i - 1][j - 1][k] + dp[i - 1][j + 1][k]

 

이 때, 범위 체크할 수는 마지막 수인 j밖에 남지 않았으므로 0일 때, 9일 때로 분기해서 확인해주면 된다.

이 부분이 초기 점화식에서 문제가 생겼던 부분이며, 비트를 도입해서 해결한 부분이다.

if (j == 0) dp[i][j][bit] = (dp[i][j][bit] + dp[i - 1][j + 1][k]) % MOD;
else if (j == 9) dp[i][j][bit] = (dp[i][j][bit] + dp[i - 1][j - 1][k]) % MOD;
else dp[i][j][bit] = (dp[i][j][bit] + dp[i - 1][j - 1][k] + dp[i - 1][j + 1][k]) % MOD;

---

초기 점화식으로도 메모리나 시간이 부족하지 않았기 때문에 비트마스킹은 생각을 못했다.

(물론 초기 점화식이 틀렸을 수도 있다.)

 

인터넷을 참고하면서 풀었는데 숫자와 관련된 조건이 많은 문제이면 비트마스킹을 한 번 떠올려라도 보는게 좋겠다.