LOJ#6503.「雅礼集训 2018 Day4」Magic[容斥+NTT+启发式合并]

题意

\(n\) 张卡牌 \(m\) 种颜色，询问有多少种本质不同的序列满足相邻颜色相同的位置数量等于 \(k\)。

分析

首先本质不同不好直接处理，可以将同种颜色的卡牌看作是不相同的，求出答案后除以 \(\prod {a_i!}\) 即可。

如果我们能够得到一个至少存在 \(k\) 个魔术对的排列数，就可以容斥了。

考虑单独处理每种颜色, 枚举一个颜色 \(i\)，计算这种颜色至少有 \(j\) 对的方案总数。

可以选择 \(j\) 张牌保证这些牌一定跟在某张牌的后面，这样就可以形成 \(\geq j\) 个满足要求的魔术对了。

然后做一个背包。

定义状态 \(f_{i,j}\) 表示处理了前 \(i\) 种颜色，至少存在 \(j\) 对魔术对的排列数。

对于 \(f_{m,j}\)，剩下的 \(n-j\) 张牌可以任意排列。背包是一个卷积的形式可以 NTT 优化。

合并背包的过程可以用堆维护，每次合并堆顶最小的两个背包数组，复杂度类似启发式合并。

所以总时间复杂度为 \(O(n\log^2n)\)。

技巧：在计算方案出现困难的时候考虑容斥简化条件(类似反演题目的思想)

代码

#include<bits/stdc++.h>
using namespace std;
#define rep(i,a,b) for(int i=a;i<=b;++i)
#define pb push_back
typedef long long LL;
typedef pair<int,int> pii;
#define mp make_pair
inline int gi(){
	int x=0,f=1;char ch=getchar();
	while(!isdigit(ch))	{if(ch=='-') f=-1;ch=getchar();}
	while(isdigit(ch)){x=(x<<3)+(x<<1)+ch-48;ch=getchar();}
	return x*f;
}
template<typename T>inline bool Max(T &a,T b){return a<b?a=b,1:0;}
template<typename T>inline bool Min(T &a,T b){return b<a?a=b,1:0;}
const int N=1e5 + 7,mod=998244353;
int n,m,k,K,len,lst,sz[N<<1],a[N];
LL invfac[N],inv[N],fac[N];
vector<LL>v[N<<1];
priority_queue<pii,vector<pii>,greater<pii> >Q;
void add(LL &a,LL b){a+=b;if(a>=mod) a-=mod;}
LL Pow(LL a,LL b){
	LL res=1ll;
	for(;b;b>>=1,a=a*a%mod) if(b&1) res=res*a%mod;
	return res;
}
LL Inv(LL a){return Pow(a,mod-2);}
LL C(int n,int m){
	return fac[n]*invfac[m]%mod*invfac[n-m]%mod;
}
void ntt(vector<LL> &c,int n,int f){
	for(int i=0,j=0;i<n;++i){
		if(i<j) swap(c[i],c[j]);
		for(int k=n>>1;(j^=k)<k;k>>=1);
	}
	for(int i=1;i<n;i<<=1){
		LL wn=Pow(3,(mod-1)/(i<<1));
		if(f==-1) wn=Inv(wn);
		for(int j=0;j<n;j+=i<<1){
			LL w=1;
			for(int k=0;k<i;++k,w=w*wn%mod){
				LL x=c[j+k],y=w*c[j+k+i]%mod;
				c[j+k]=(x+y)%mod;
				c[j+k+i]=(x-y+mod)%mod;
			}
		}
	}
	if(f==-1){LL inv=Inv(n);for(int i=0;i<n;++i) c[i]=c[i]*inv%mod;}
}
int main(){
	m=gi(),n=gi(),K=gi();
	fac[0]=invfac[0]=inv[1]=1;
	rep(i,1,n){
		if(i^1) inv[i]=(mod-mod/i)*inv[mod%i]%mod;
		fac[i]=fac[i-1]*i%mod;
		invfac[i]=invfac[i-1]*inv[i]%mod;
	}
	rep(c,1,m){
		sz[c]=a[c]=gi();int x=a[c];
		Q.push(mp(sz[c],c));
		v[c].resize(x+7);
		rep(i,0,x-1) v[c][i]=C(x,i)*fac[x-1]%mod*invfac[x-i-1]%mod;
	}
	while(!Q.empty()){
		int x=Q.top().second;Q.pop();
		if(Q.empty()) { lst=x; break; }
		int y=Q.top().second;Q.pop();
		for(len=1;len<=sz[x]+sz[y];len<<=1);
		v[x].resize(len+7);
		v[y].resize(len+7);
		ntt(v[x],len,1);
		ntt(v[y],len,1);
		for(int i=0;i<len;++i) v[x][i]=v[x][i]*v[y][i]%mod;
		ntt(v[x],len,-1);

		sz[x]+=sz[y];
		v[y].clear();
		Q.push(mp(sz[x],x));
	}
	LL ans=0;
	rep(i,K,n-1)
		add(ans,(i-K&1?mod-1:1)*v[lst][i]%mod*C(i,K)%mod*fac[n-i]%mod);
	rep(i,1,m) ans=ans*invfac[a[i]]%mod;
	printf("%lld\n",ans);
	return 0;
}