AT_arc170_c Prefix Mex Sequence

原题链接

发现所有的限制都只跟已经使用的数字个数有关，故记 \(f_{i,j}\) 为前 \(i\) 个数，使用了 \(j\) 个数的方案数。

当 \(s_i=0\)，可以填已经出现过的数，也可以填没有出现过但并非当前 \(\operatorname{mex}\) 的数，即 \(f_{i,j}=f_{i-1,j}\times j +f_{i-1,j-1}\times (m-(j-1))\)。

当 \(s_i=1\)，只能填 \(\operatorname{mex}\) 所以 \(f_{i,j}=f_{i-1,j-1}\)。

最后答案是 \(\sum\limits_{i=0}^{\min(n,m+1)} f_{n,i}\)

#include<stdio.h>
#include<iostream>
#include<algorithm>
#define int long long
using namespace std;
const int MAXN=5e3+10,MOD=998244353;
int n,m,s[MAXN],f[MAXN][MAXN],now,ans;
signed main()
{
    cin.tie(0),cout.tie(0);
    ios::sync_with_stdio(0);
    cin>>n>>m;f[0][0]=1;
    for(int i=1;i<=n;++i)
    {
        cin>>s[i];
        if(!s[i])
        {
            for(int j=0;j<=n;++j)
            {
                f[i][j]=(f[i][j]+f[i-1][j]*j%MOD)%MOD;
                if(m-j>=0)
                    f[i][j+1]=f[i-1][j]*(m-j)%MOD;
            }
        }
        if(s[i])
        {
            for(int j=0;j<=n;++j)
                f[i][j+1]=f[i-1][j];
            ++now;
        }
    }
    for(int j=0;j<=min(n,m+1);++j)
        ans=(ans+f[n][j])%MOD;
    cout<<ans<<'\n';return 0;
}