【ABC215】E - Chain Contestant

atcoder.jp

解法

問題を簡単に言い換えると、同じ種類のコンテストには連続でしか出場できないという条件を満たす、コンテストへの出場の仕方は何通りですか？
といった感じになる。

DPで解くことを考えてみる。

どの種類のコンテストへの出場経験があるか、という状態が欲しくなる。
最大10種類のコンテストが開催される可能性があるので、出場経験としてありえる組み合わせは2¹⁰通りになる。
既に出場したコンテストの集合を状態として持つと考えると、bitDPで実現できる。

ただし、同じ種類のコンテストなら連続で出場できることため、前回どのコンテストに出場したかという状態も必要になる。

よって以下のようなDPで解くことができる。
dp[i][S][k] ← i回目までコンテストが開催されたときに、集合Sに含まれるコンテストへの出場経験があり、前回出場したコンテストがk種類目である時が何通りか

遷移

i番目のコンテストに出場しない場合

出場経験も前回出場したコンテストも変化しません。

dp[i+1][S][k] += dp[i][S][k]

i番目のコンテストに出場する場合

i番目のコンテストに出場するには下の条件を満たす必要があります。

• i番目のコンテストと同じ種類のコンテストへの出場経験がない
• i番目のコンテストと前回出場したコンテストの種類が同じ

i番目のコンテストの種類がv番目だった場合、下のように遷移します。

dp[i+1][S ∪ v][v+1] += dp[i][j][k]

答え

N番目のコンテストの遷移が終了した時点で、ありえる出場経験の総和が答えになります。

実装

void solve(long long N, std::string S) {
    vector<vector<vector<mint>>> dp(N+1, vector<vector<mint>>(3000, vector<mint>(11, 0)));
    dp[0][0][0] = 1;

    REP (i, N) {
        int v = S[i] - 'A';
        REP (j, 1025) {
            REP (k, 11) {
                // i番目のコンテストに出場する場合
                if (!(j & (1 << v)) || (v+1) == k) {
                    dp[i+1][j | (1<<v)][v+1] += dp[i][j][k];
                }
                // 出場しない場合
                dp[i+1][j][k] += dp[i][j][k];
            }
        }
    }

    // 答えを計算
    mint ans = 0;
    for (int i = 1; i < 1025; i++) {
        REP (k, 11) {
            ans += dp[N][i][k];
        }
    }

    c(ans); // 出力
}

感想

方針はすぐに立ったけど実装に手間取ってしまった。
ビット演算を素早く実装できるようになりたい。