【AtCoder】ARC170 C - Prefix Mex Sequence

問題

解法

「mex を追加する、または mex 以外を追加する」という操作において、種類数という情報だけで遷移元と遷移先を表現できる（情報として十分である）ということに注目する。

dp[i] = 数列の要素が i 種類 とすると

• mex を追加する場合: 種類数だけで考えたら 1 種類増えるだけ（増えて M + 1 種類を超えることはないが）

• mex を追加しない場合: 1 種類増えたりしなかったりするが常に mex 以外は M 種類なので遷移先がわかれば係数はわかる

M と N の大小関係に気を使わないと TLE してしまうため、新しく種類数の上限を K として求め、0 から K まで順番に見ていきながら適宜ダメな遷移先をカットするイメージで実装すると良い。

#include "my_template.hpp"
#include "math/static_modint.hpp"
using namespace std;

using mint = mint998;

int main() {
    INT(N, M);
    VEC(int, S, N);
    const int K = min(M + 1, N);  // 種類数の上限
    vector<mint> dp(K + 1, 0);    // dp[i] = 数列の要素が i 種類
    dp[0] = 1;
    REP(i, N) {
        vector<mint> np(K + 1, 0);
        if (S[i] == 1) {
            // mex を入れるので必ず j 種類 -> (j + 1) 種類になる
            REP(j, K + 1) {
                if (j + 1 > M + 1) break;
                if (j + 1 <= K) np[j + 1] += dp[j];
            }
        }
        if (S[i] == 0) {
            // (M + 1) 種類のうち, mex にならない M 種類のどれかを追加
            REP(j, K + 1) {
                // 種類数が増える
                if (j + 1 <= K) np[j + 1] += dp[j] * (M - j);
                // 種類数が増えない
                np[j] += dp[j] * j;
            }
        }
        swap(dp, np);
    }
    mint ans = SUM(dp, mint(0));
    print(ans);
    return 0;
}

感想

種類数に注目すれば十分という発想がなかった。

である要素数に注目し、適宜 に対応する部分に値を埋めていく方法を考えていたが、埋めた要素の組合せによっては後に破綻することがある。

状態を細かく持てばできるが、状態をまとめることができなかった。

種類数が一種の不変量（？）（その情報だけで閉じているみたいな）になっていると見ることもできて、そういう意味では（広義の）不変量の発見力が足りないとも言える。