公式

L - Mex on Blackboard 2 解説 by Chippppp


まず、各々の操作において、追加できる数は \([0, mex(A)]\) であることがわかります。

追加できる数は \(mex(A)\) のみに依存するので、次の動的計画法を考えることができます:

\(dp_{i, j} = (i回操作したとき mex(A) = j となるAの個数)\)

よって \(O(N(N + K))\) で答えを求めることができます。

(新しい数が追加されたとき \(mex(A)\) が何に変化するかを事前に計算しておく必要があります。)

#include <bits/stdc++.h>
#include <atcoder/all>

int main() {
    using namespace std;
    using mint = atcoder::modint998244353;

    int n, k;
    cin >> n >> k;
    vector<int> a(n);
    for (auto& i: a) cin >> i;

    unordered_set<int> st(a.begin(), a.end());
    int mex = 0;
    while (st.find(mex) != st.end()) ++mex;

    // m[i] = (iがはじめて追加された時のmex(A))
    vector<int> m(4001);
    int i = mex;
    while (i < 4001) {
        int j = i + 1;
        while (st.find(j) != st.end()) ++j;
        m[i] = j;
        i = j;
    }

    // DP
    vector dp(k + 1, vector(4002, mint(0)));
    dp[0][mex] = 1;
    for (int i = 0; i < k; ++i) for (int j = 0; j < 4001; ++j) {
        dp[i + 1][j] += dp[i][j] * j;
        dp[i + 1][m[j]] += dp[i][j];
    }

    // dp.back()の和を出力
    cout << reduce(dp.back().begin(), dp.back().end()).val() << '\n';
}

投稿日時:
最終更新: