题解：AT_arc170_c [ARC170C] Prefix Mex Sequence

洛谷传送门

AT传送门

题意

给定一个长度为 \(N\) 的序列 \(S\)，每个元素为 \(0\) 或 \(1\)。给定一个数 \(m\)，要求构造序列 \(A\)，使得 \(\forall1\le i \le n\)，有 \(A_i \in [0,m]\)，且当且仅当 \(S_i=1\)时，\(A_i=\mathrm{mex}((A_1,A_2,\dots A_{i-1}))\)。

思路

我们不考虑 \(\mathrm{mex}\) 的具体值，只考虑在它之前选了多少个数。

当 \(n \le m+1\) 时，不论我们怎么选，每一位选择时总有数没被我们选过。遍历到点 \(i\) 时，若 \(S_i=1\)，则只能选择前缀 \(\mathrm{mex}\) 值，只有一种选法。同理，若 \(S_i=0\)，有 \(m\) 种选法。

若 \(n> m+1\)，则我们设计 dp 状态 \(dp_{i,j}\) 表示点 \(i\) 时区间 \([1,i]\) 选择了 \(j\) 个数。若 \(S_i=1\) 时，只能选 \(\mathrm{mex}\)，\(dp_{i,j}=dp_{i-1,j-1}\)。当 \(S_i=0\)，可以从已经选择的数值中选一个，或者从未选择的数中选择一个不为 \(\mathrm{mex}\) 的数。因此，有 \(dp_{i,j}=(m+1-j)dp_{i-1,j-1}+j dp_{i-1,j}\)。

初始状态 \(dp_{0,0}=1\)。最终答案为 \(\sum _{i=1}^{m+1}dp_{n,i}\)。

时间复杂度 \(O(n\min(n,m))\)。

代码

#include <bits/stdc++.h>
using namespace std;
#define INF 0x3f3f3f3f3f3f3f3fll
#define IINF 0x3f3f3f3f
#define DINF 10000000
#define ll long long
#define sc scanf
#define pr printf
#define v1 first
#define v2 second
#define lowbit(x) ((x)&(-x))
const int N=5e3+5;
ll dp[N][N];
int n,m;
int s[N];
const ll mod=998244353;
int main() {
    sc("%d%d",&n,&m);
    for (int i=1; i <= n; i++)
        sc("%d",&s[i]);
    m++;
    if (n<=m) {
        ll ans=1;
        for (int i=1; i <= n; i++)
            if (!s[i])ans=ans*(m-1)%mod;
        return 0;
    }
    dp[0][0]=1;
    for (int i=1; i <= n; i++) {
        for (int j=1; j <= m; j++) {
            if (s[i])
                dp[i][j]=dp[i-1][j-1];
            else {
                dp[i][j]+=dp[i-1][j-1]*(m-j)%mod;
                dp[i][j]%=mod;
                dp[i][j]+=dp[i-1][j]*j%mod;
                dp[i][j]%=mod;
            }
        }
    }
    ll ans=0;
    for (int i=1;i <= m; i++)
        ans=(ans+dp[n][i])%mod;
    cout << ans;
    return 0;
}