题解：[ARC170C] Prefix Mex Sequence

题目传送门

题意分析

对于 \(s_i=1\) 的 \(a_i\)，显然 \(a_i\) 时由 \(a_1,a_2,\cdots,a_{i-1}\) 确定的唯一的数。

当 \(n\leq m\) 时，最大的 \(\operatorname{mex}\) 的值一定不会超过 \(m\)。故此时所有 \(s_i=0\) 的 \(a_i\) 都可以在 \(\set{0,1,\cdots,m}\setminus\set{\operatorname{mex}(\left\langle a_1,a_2,\cdots,a_{i-1}\right\rangle)}\) 中任意选取，选取方案数为 \(m\)。

故，\(n\leq m\) 时，记 \(\textit{cnt}_0\) 为 \(s_1,s_2,\cdots,s_n\) 中 \(0\) 的个数，答案为 \(m^{\textit{cnt}_0}\)。

---

现在考虑 \(n>m\) 时。考虑到这时不能再任意选取，因为有可能某些情况下后面是无解的（找不到一个 \(\operatorname{mex}\) 的值来满足 \(s_i=1\) 的需求），但是同样的 \(s\) 可能前面的 \(a_i\) 的值更改一下，后面就又有解了。这样不能直接计算，考虑 DP。

设 \(\textit{dp}_{i,j}\) 表示处理了 \(a_1\sim a_i\)，\(a_1,a_2,\cdots,a_i\) 中有 \(j\) 个不同的数的方案数。

• 若 \(s_i=1\)，则 \(a_i\) 确定，因为 \(a_i\) 是 \(a_1,a_2,\cdots,a_{i-1}\) 中未出现的数，因此不同的数的个数会增加，故有：

  \[\textit{dp}_{i,j}=\textit{dp}_{i-1,j-1} \]

• 否则若 \(s_i=0\)，则：

  • \(a_i\) 是之前出现过的数，方案数为 \(j\cdot \textit{dp}_{i-1,j}\)。即枚举每一个可能的 \(a_i=j\)。
  • \(a_i\) 是之前没有出现过的数，方案数为 \((m+1-j)\textit{dp}_{i-1,j-1}\)。因为在 \(a_{i-1}\) 时 \(\set{0,1,\cdots,m}\) 中已经出现过的数的数量为 \(j-1\)，并且还要去除 \(\operatorname{mex}(\langle a_1,a_2,\cdots,a_{i-1}\rangle)\)。

  故有：

  \[\textit{dp}_{i,j}=j\cdot \textit{dp}_{i-1,j-1}+(m+1-j)\textit{dp}_{i-1,j-1} \]

\(\mathcal O(nm)\) DP 求解即可。

AC 代码

//#include<bits/stdc++.h>
#include<algorithm> 
#include<iostream>
#include<cstring>
#include<iomanip>
#include<cstdio>
#include<string>
#include<vector>
#include<cmath>
#include<ctime>
#include<deque>
#include<queue>
#include<stack>
#include<list>
using namespace std;
constexpr const int N=5000,M=1e9,P=998244353;
int n,m,s[N+1];
namespace judge1{
	int qpow(int base,int n){
		int ans=1;
		while(n){
			if(n&1){
				ans=1ll*ans*base%P;
			}
			base=1ll*base*base%P;
			n>>=1;
		}
		return ans;
	}
	int main(){
		int cnt=0;
		for(int i=1;i<=n;i++){
			cnt+=!s[i];
		}
		cout<<qpow(m,cnt)<<'\n';
		return 0;
	}
} 
namespace judge2{
	int dp[N+1][N+1];
	int main(){
		dp[0][0]=1;
		for(int i=1;i<=n;i++){
			for(int j=1;j<=n;j++){
				if(s[i]){
					dp[i][j]=dp[i-1][j-1];
				}else{
					dp[i][j]=(1ll*j*dp[i-1][j]%P+(m-j+1ll)*dp[i-1][j-1]%P)%P;
				}
			}
		}
		int ans=0;
		for(int i=1;i<=m+1;i++){
			ans=(ans+dp[n][i])%P; 
		}
		cout<<ans<<'\n';
		return 0;
	}
}
int main(){
	/*freopen("test.in","r",stdin);
	freopen("test.out","w",stdout);*/
	
	ios::sync_with_stdio(false);
	cin.tie(0);cout.tie(0);
	
	cin>>n>>m;
	for(int i=1;i<=n;i++){
		cin>>s[i];
	}
	if(n<=m){
		judge1::main();
	}else{
		judge2::main();
	}
	
	cout.flush();
	
	/*fclose(stdin);
	fclose(stdout);*/
	return 0;
}