AtCoder Regular Contest 058 E - 和風いろはちゃん / Iroha and Haiku

問題

arc058.contest.atcoder.jp

解法

余事象を考えます。つまり, 「XYZ が含まれない数列の数」を求めることを目的にします。これを dp で解くことを考えましょう。

状態として, 直前の X+Y+Z 個の数字を持っておけば良いのはわかりますが, これだと 10^(X+Y+Z) になって間に合いません。ですが, 状態としては 合計が X+Y+Z 以下である直前の数字だけ見ていれば良いことがわかります(例えば X+Y+Z = 6 のとき, 直前から順番に(1, 2, 3) となってるのは覚えておく必要があるけど (1, 2, 3, 2) とかまでは覚えておく必要がない)。

ここからは実装の話です。上記の状態をどう表現するかですが, これは bit で表現することができます。1 を "1", 2 を "01", 3 を "001", 4 を "0001" というようにします。

これで肝心なのは, 「区切りとしてあり得るものが 1 で表現されている」という点です。例えば, 直前から順番に (1, 2, 3) というのは "1101001" (最初の 1 は「0 番目に区切りがある」ということを示す)と表現されますが, これで bit が立っているところを見ていくと 0, 1, 3, 6 番目の bit が立っており, 確かに区切りが 1, 3, 6 になっていることが分かります。

また, 状態 s に i を付け加えるというのは, 「s の区切りで j 番目に区切りがついていたのは i+j 番目に区切りを入れるようにする」ということと「0 番目に区切りを入れる」ということをやれば良いので, (s < < i)+1 で表現できます(X+Y+Z 個分の状態しかいらないので mask をかける必要はあります)。

後は「今の状態が NG 状態でないか」ということですが, dp の遷移をやる際は「s では NG ではない」という前提があるのでそれは考えることにすると, 数を追加した際に NG になるには, 「追加して初めて Z, Z+Y, Z+Y+X に区切りがついた(Z, Z+Y, Z+Y+X の位置に bit が立っている)」ということが必要十分です。

const int MOD = 1e9+7;

ll dp[44][1<<18];

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int N, X, Y, Z;
    cin >> N >> X >> Y >> Z;
    ll ans = 1;
    for (int i = 0; i < N; i++) {
        ans *= 10;
        ans %= MOD;
    }
    int ng = 0;
    ng |= 1<<Z;
    ng |= 1<<(Y+Z);
    ng |= 1<<(X+Y+Z);
    int mask = 1<<(X+Y+Z+1); mask--;
    dp[0][1] = 1;
    for (int n = 0; n < N; n++) for (int state = 0; state <= mask; state++) {
        for (int i = 1; i <= 10; i++) {
            int ns = ((state<<i)&mask)+1;
            if ((ns&ng) == ng) continue;
            dp[n+1][ns] += dp[n][state];
            dp[n+1][ns] %= MOD;
        }
    }
    for (int i = 0; i < mask; i++) {
        ans -= dp[N][i];
        ans %= MOD;
    }
    cout << (ans+MOD)%MOD << endl;
    return 0;
}

「合計が S になるまでの数の並びの状態」を bit で表現するのは頭良い。自分もだいたい方針あってたけど bit で表現してなかったので TLE した。