CF1608F MEX counting 解题总结

\(\small \mathsf{date:20260920}\)

闲话

燃尽了\(\cdots\cdots\) 先溜过来写一点总结。

黑题还是太难搞了

题面

对于一个非负整数数组 \(c\)，\(MEX(c)\) 表示不在数组中出现的最小非负整数。例如，\(MEX([0, 1, 3]) = 2\)，\(MEX([42]) = 0\)。

给定整数 \(n, k\)，以及一个数组 \([b_1, b_2, \ldots, b_n]\)。

请你求满足以下条件的数组 \([a_1, a_2, \ldots, a_n]\) 的个数：

• 对于每个 \(i\)，\(0 \le a_i \le n\)。
• 对于每个 \(i\)（\(1 \le i \le n\)），都有 \(|MEX([a_1, a_2, \ldots, a_i]) - b_i| \le k\)。

由于答案可能很大，请输出对 \(998\,244\,353\) 取模后的结果。

MEX？

已经见过 \(INF\) 次这个东西了，感觉极其恶心，因此这里先拿出来说一下。

\(MEX\) 表示不出现在数组中的最小非负整数。在非负整数定义下：

• \(MEX\) 数组具有集合的性质，只关心元素是否出现，并不关心顺序和数量。
• 若 \(MEX(c)=x\)，则 \(0,1,\dots ,x-1\) 均在 \(c\) 中。
• 若 \(MEX(c)=x\)，则大于 \(x\) 的元素是否出现对 \(MEX\) 值没有影响。
• 若往 \(MEX\) 数组中添加元素 \(x\)，则:
  • 无论添加什么，\(MEX\) 的值都不会减小，即其具有单调性。
  • 当 \(x<MEX(c) 或 x>MEX(c)\)，\(MEX\) 的值不变。
  • 当 \(x=MEX(c)\)，\(MEX\) 的值会“跳跃式”地增加——跳过所有已出现过的连续整数，停在第一个没出现过的位置。
  • 换言之，添加的大于 \(MEX\) 的值只是暂时无用，它的影响会在下一次添加 \(x=MEX\) 时清算。

显然，\(MEX\) 的性质非常适合“贡献挂起”的DP。这对我们解这道题有很大的帮助。

解题思路

和前面两道题一样，这道题需要用到暂时挂起的思想。

设 \(dp_{i,j,k}\) 表示当前处理到前 \(i\) 个数，前 \(i\) 个数的的 \(MEX\) 为 \(j\)，挂起了 \(k\) 个大于 \(MEX\) 的数。

枚举 \(i,j,k\)，对于第 \(i\) 个数，我们可以：

• 若它小于 \(j\)，没有什么影响，\(dp_{i+1,j,k}+=dp_{i,j,k}\times j\)

• 若它大于 \(j\) ：

  • 若它已经被挂在了序列中，没有什么影响，\(dp_{i+1,j,k}+=dp_{i,j,k}\times k\)
  • 若它没有被挂在序列中，挂起它，\(dp_{i+1,j,k+1}+=dp_{i,j,k}\)。注意这里没有系数，是因为挂起造成的方案贡献最后统一清算。

• 若它等于 \(j\)，枚举所有的下一个 \(MEX\)，加上其方案数，\(dp_{i+1,newj,newk}+=dp_{i,j,k}\times \binom{k}{newj-j-1}\times (newj-j-1)!\)

  细节

  对于所有的枚举，注意到 \(|j - b_i| \le k\)（无论是 \(j\) 还是 \(newj\) 都算），因此每次最后要将不合法的状态置零（可以利用这个剪枝）

最后计算时，对于所有挂起的 \(k\) 我们都可以将其任意排列（因为在DP时我们将所有挂起的方案视为等价），因此每一项要乘上 \(A_{n-j}^{k}\)

40pts代码

#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N=2005,K=55,mod=998244353;
int dp[2][N][N],n,kk,b[N],fac[N]={1},facinv[N]={1},now=0;
int qpow(int a,int b,int mod)
{
    int res=1%mod;
    a%=mod;
    while(b)
    {
        if(b&1) res=res*a%mod;
        a=a*a%mod;
        b>>=1;
    }
    return res;
}
int inv(int a)
{
    return qpow(a,mod-2,mod);
}
int getC(int a,int b)
{
    return fac[a]*facinv[a-b]%mod*facinv[b]%mod;
}
signed main()
{
    cin>>n>>kk;
    for(int i=1;i<=n;i++)cin>>b[i];
    for(int i=1;i<N;i++)fac[i]=fac[i-1]*i%mod,facinv[i]=inv(fac[i]);
    dp[0][0][0]=1;
    for(int i=0;i<n;i++)
    {
        memset(dp[now^1],0,sizeof dp[now^1]);
        for(int j=0;j<=n;j++)
        {
            for(int k=0;k<=n;k++)
            {
                for(int newj=max(j+1,b[i+1]-kk);newj<=min(n,b[i+1]+kk);newj++)
                {
                    int newk=k-newj+j+1;
                    if(newk<0)continue;
                    if(abs(newj-b[i+1])>kk)continue;
                    dp[now^1][newj][newk]=(dp[now^1][newj][newk]+dp[now][j][k]*getC(k,newj-j-1)%mod*fac[newj-j-1]%mod+mod)%mod;
                }
                dp[now^1][j][k]=(dp[now^1][j][k]+(j+k)*dp[now][j][k]%mod+mod)%mod;
                dp[now^1][j][k+1]=(dp[now^1][j][k+1]+dp[now][j][k]%mod+mod)%mod;
            }
        }

        for(int j=0;j<=n;j++)
            if(j<max(0LL,b[i+1]-kk)||j>min(n,b[i+1]+kk))
                memset(dp[now^1][j],0,sizeof dp[now^1][j]);
        now^=1;
    }
    int ans=0;
    for(int j=0;j<=n;j++)
    {
        for(int k=0;k<=n;k++)
        {
            ans=(ans+dp[now][j][k]*fac[n-j]%mod*facinv[n-j-k]%mod+mod)%mod;
        }
    }
    cout<<ans;
    return 0;
}

优化

刚刚的做法是 \(O(n^3k)\) 的，显然无法通过。

回顾转移方程：

\(dp_{i+1,newj,newk}+=dp_{i,j,k}\times \binom{k}{newj-j-1}\times (newj-j-1)!\)

我们可以将其写成这样的形式：

设 \(d=newj-j-1,newk=k-d\)，则

\(dp_{i+1,newj,newk}+=dp_{i,j,k}\times \binom{k}{newj-j-1}\times (newj-j-1)!\)

可以改写为

$ dp_{i+1,newj,newk}+=dp_{i,j,k}\times \binom k d \times d !$

展开，得

$\binom k d \times d！=\frac{k!}{d!(k-d)!}\times d!=\frac{k!}{(k-d)!} $

所以原式等于

\(dp_{i+1,newj,newk}+=dp_{i,j,newk+newj-j-1}\times\frac{(newk+newj-j-1)!}{newk!}\)

下标中又有 \(newk\) 又有 \(newj\) ，似乎有些麻烦，想办法处理掉：

定义 \(g_{i,j,s}=dp_{i,j,k}\cdot k!\)，其中 \(s=j+k\)

则 \(dp_{i,j,k}=\frac{g_{i,j,j+k}}{k!}\)

代入原式，\(dp_{i+1,newj,newk}+=\frac{g_{i,j,newk+newj-1}}{(newk+newj-j-1)!}\cdot\frac{(newk+newj-j-1)!}{newk!}\)

于是，阶乘正好约掉！

\(dp_{i+1,newj,newk}+=\frac{g_{i,j,newk+newj-1}}{newk!}\)

写成前缀和，砍掉一层循环

\(dp_{i+1,newj,newk}+=\frac{1}{newk!}\sum_{j=0}^{newj-1}g_{i,j,newk+newj-1}\)

定义 \(pre_{i,j,x}=\sum_{j=0}^{newj-1}g_{i,j,x}\)

则 \(dp_{i+1,newj,newk}+=\frac{pre_{i,j,x}}{newk!}\)，取余后 \(pre_{i,j,x}\times facinv_{newk}\)

优化后的代码

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
const int N=2005,mod=998244353;
static unsigned dp[2][N][N];
ll fac[N],ifac[N],pre[N];
int n,kk,b[N];
ll pw(ll a,ll e)
{
	ll r=1;
	a%=mod;
	while(e)
	{
		if(e&1) r=r*a%mod;
		a=a*a%mod;
		e>>=1;
	}
	return r;
}
int main()
{
	scanf("%d %d",&n,&kk);
	for(int i=1;i<=n;i++) scanf("%d",&b[i]);
	fac[0]=1;
	for(int i=1;i<N;i++) fac[i]=fac[i-1]*i%mod;
	ifac[N-1]=pw(fac[N-1],mod-2);
	for(int i=N-1;i>0;i--) ifac[i-1]=ifac[i]*i%mod;
	int cur=0;
	dp[cur][0][0]=1;
	int lo0=0,hi0=0;
	for(int i=1;i<=n;i++)
	{
		int lo=max(0,b[i]-kk),hi=min(n,b[i]+kk);
		if(lo>hi)
		{
			puts("0");
			return 0;
		}
		cur^=1;
		for(int j=lo;j<=hi;j++) memset(dp[cur][j],0,sizeof(unsigned)*(n+2));
		int smax=i-1;
		for(int j=max(lo,lo0);j<=min(hi,hi0);j++)
			for(int s=j;s<=smax;s++)
			{
				ull v=dp[cur^1][j][s];
				if(!v) continue;
				dp[cur][j][s]=(dp[cur][j][s]+v*s)%mod;
				dp[cur][j][s+1]=(dp[cur][j][s+1]+v)%mod;
			}
		for(int s=0;s<=smax;s++) pre[s]=0;
		for(int jp=lo0+1;jp<=hi;jp++)
		{
			int j=jp-1;
			if(j>=lo0&&j<=hi0)
				for(int s=j;s<=smax;s++)
					pre[s]=(pre[s]+(ull)dp[cur^1][j][s]*fac[s-j])%mod;
			if(jp>=lo)
				for(int s=jp-1;s<=smax;s++)
					dp[cur][jp][s+1]=(dp[cur][jp][s+1]+pre[s]*ifac[s+1-jp])%mod;
		}
		lo0=lo;
		hi0=hi;
	}
	ll ans=0;
	for(int j=lo0;j<=hi0;j++)
		for(int s=j;s<=n;s++)
			ans=(ans+(ull)dp[cur][j][s]*fac[n-j]%mod*ifac[n-s])%mod;
	printf("%lld\n",ans);
}