JOI 2008/2009 予選 6 「ビンゴ」解法

後輩への解説用に作ったものをブログにも上げておきます．誰かの役に立ちますように．いろいろ独断と偏見が含まれているので，上級者の方は温かい目でご覧ください．

原本：https://drive.google.com/open?id=1pa_I0VqJWKSWOIY6jMg8Cky8BY649gej

STEP1 普通に DP

添え字として，今どのマスを見ているか，1 つ前のマスの値は何であったか，1 つ前のマスまでの合計値はいくつか，というデータを持つ．遷移として，1 つ前のマスの値より大きいものを全探索する．時間計算量は O(N²M²S)．

サンプルコード

int N, M, S;
// now := N*Nマス中のどこか
// bef := 1つ前のマスの値+1 (今回はbef以上)
// sum := 1つ前のマスまでの合計値
int memo[50][2001][3001];
int solve(int now, int bef, int sum){
    if(sum > S) return 0;
    if(now == N*N){
        if(sum == S) return 1;
        else return 0;
    }
    if(memo[now][bef][sum] != -1) return memo[now][bef][sum];
    int ans = 0;
    for(int i=bef; i<=M; i++){
        ans += solve(now+1, i+1, sum+i);
    }
    return memo[now][bef][sum] = ans % MOD;
}

int main(){
    cin >> N >> M >> S;
    for(int i=0; i<N*N; i++){
        for(int j=0; j<=M; j++){
            for(int k=0; k<=S; k++){
                memo[i][j][k] = -1;
            }
        }
    }
    cout << solve(0, 1, 0) << endl;
    return 0;
}

STEP2 遷移をまとめる

DP の時間計算量を落とすためにできることは，添え字を減らすか，solve 関数の中のループを無くすかのどちらかである．ここでは solve 関数の中のループを無くせる．具体的に考えてみる．

ある関数呼び出し solve(i, j, k) が次に呼び出すのは solve(i+1, j, k+j), solve(i+1, j+1, k+j+1), solve(i+1, j+2, k+j+2), …, solve(i+1, M, k+M) である．また，ある関数呼び出し solve(i, j+1, k) が次に呼び出すのは，solve(i+1, j+1, k+j+1), solve(i+1, j+2, k+j+2), …, solve(i+1, M, k+M) である．下線部において，呼び出す遷移先が共通しているから，うまくまとめることが出来る．今回の場合，例えば関数呼び出し solve(i, j, k) が次に呼び出す関数を，solve(i+1, j, k+j), solve(i, j+1, k) とすればよい．

時間計算量は O(N²MS)．

サンプルコード

int N, M, S;
// now := N*Nマス中のどこか
// bef := 1つ前のマスの値+1 (今回はbef以上)
// sum := 1つ前のマスまでの合計値
int memo[50][2001][3001];
int solve(int now, int bef, int sum){
    if(sum > S) return 0;
    if(now == N*N){
        if(sum == S) return 1;
        else return 0;
    }
    if(bef > M) return 0;
    if(memo[now][bef][sum] != -1) return memo[now][bef][sum];
    int ans = solve(now+1, bef+1, sum+bef) + solve(now, bef+1, sum);
    return memo[now][bef][sum] = ans % MOD;
}

int main(){
    cin >> N >> M >> S;
    for(int i=0; i<N*N; i++){
        for(int j=0; j<=M; j++){
            for(int k=0; k<=S; k++){
                memo[i][j][k] = -1;
            }
        }
    }
    cout << solve(0, 1, 0) << endl;
    return 0;
}

STEP3 配列を使いまわす

この問題の本質的な改善は STEP2 だが，まだ使用メモリ量が多い．メモ化再帰ではなく for ループを用いる形に DP を書き換えることで，配列を使いまわすことが可能な場合がある．書き換えの際には，再帰の末端の方から先に処理を行うことを意識する．今回の場合，now と bef は小さくなるようループを回さなければならない．sum はどちらでもよい．空間計算量は O(MS)．

「配列を使いまわす」とは，配列の特定の添え字の要素数を 2 にして，呼び出し時には %2 とし，同じメモリに別の意味を持たせながら繰り返し使う典型テクニックである．

サンプルコード

int N, M, S;
int dp[2][2002][3001];

int main(){
    cin >> N >> M >> S;
    // now := N*Nマス中のどこか
    // bef := 1つ前のマスの値+1 (今回はbef以上)
    // sum := 1つ前のマスまでの合計値
    for(int now=N*N; now>=0; now--){
        for(int bef=M+1; bef>=0; bef--){
            for(int sum=0; sum<=S; sum++){
                if(now == N*N){
                    if(sum == S) dp[now%2][bef][sum] = 1;
                    else dp[now%2][bef][sum] = 0;
                    continue;
                }
                if(bef > M){
                    dp[now%2][bef][sum] = 0;
                    continue;
                }
                int ans = (sum+bef <= S ? dp[(now+1)%2][bef+1][sum+bef] : 0) + dp[now%2][bef+1][sum];
                dp[now%2][bef][sum] = ans % MOD;
            }
        }
    }
    cout << dp[0][1][0] << endl;
    return 0;
}