[HAOI2018][bzoj5306] 染色 [容斥原理+NTT]

题面

传送门

思路

这道题的核心在于“恰好有$k$种颜色占了恰好$s$个格子”

这些“恰好”，引导我们去思考，怎么求出总的方案数呢？

分开考虑

考虑把恰好有$s$个格子的颜色，和不是$s$个颜色的格子分开来考虑

那么，显然答案可以用这样的一个式子表示：

令$lim=min(\lfloor\frac ns\rfloor,m)$，那么：

$ans=\sum_{i=0}^{lim}w_iC_miC_n^{{is}\frac{(is)!}{(s!)}i}g(m-i,n-is)$

其中$g(i,j)$表示把$i$种颜色放进$j$个格子里面，没有“恰好占了$s$个格子”的颜色的方案总数

这个$g$显然可以容斥原理来做

容斥原理

考虑对$g$用容斥原理推导，可得到如下公式：

$g(i,j)=\sum_{k=0}^i(-1)kC_i^kC_j{ks}\frac{(ks)!}{(s!)^k}(i-k){j-ks}$

把这个公式代回原式，可得到：

$ans=\sum_{i=0}^{lim}w_iC_miC_n^{{is}\frac{(is)!}{(s!)}i}\sum_{j=0}^{lim-i}(-1)jC_{m-i}^jC_{n-is}{js}\frac{(js)!}{(s!)^j}(m-i-j){n-is-js}$

我们发现这个玩意儿好像比较复杂，没有什么特殊性质

难道这条路走不通？

构造卷积

我们发现后面的这个代表$g$的这一部分里面，有很多的$j$在，

那么我们想一想，$j$的枚举是所有比当前的$i$小的，那么只要把$j$用另一个东西：$j-i$替代了，然后把$j$变成所有比当前的大的，那么是不是和这个式子等价了呢？

说干就干！

$ans=\sum_{i=0}^{lim}w_iC_miC_n^{{is}\frac{(is)!}{(s!)}i}\sum_{j=lim-i}^{lim}(-1){j-i}C_{m-i}^{{j-i}C_{n-is}}{js-is}\frac{(js-is)!}{(s!)^{j-i}}(m-j){n-js}$

诶，这样一变化......后面的式子里这么多$j-i$，前面又是$i$，这让我们想到了什么？

处理卷积啊！

变换枚举方式

我们先把这些个烦人的组合数拆成阶乘的形式

$ans=\sum_{i=0}^{{lim}w_i\frac{m!n!}{i!(is)!(m-i)!(n-is)!}\frac{(is)!}{(s!)}i}\sum_{j=lim-i}^{lim}(-1){j-i}\frac{(m-i)!(n-is)!}{(j-i)!(js-is)!(m-j)!(n-js)!}\frac{(js-is)!}{(s!)^{j-i}}(m-j){n-js}$

我们把左边的项挪进右边的最后一个sigma里面，然后把$ij$的枚举顺序反过来，得到：

$ans=\sum_{j=0}^{{lim}\sum_{i=0}}{j}w_i\frac{m!n!}{i!(is)!(m-i)!(n-is)!}\frac{(is)!}{(s!)^i}(-1){j-i}\frac{(m-i)!(n-is)!}{(j-i)!(js-is)!(m-j)!(n-js)!}\frac{(js-is)!}{(s!)^{j-i}}(m-j){n-js}$

发现有很多项可以消掉了

$ans=\sum_{j=0}^{{lim}\sum_{i=0}}{j}w_i\frac{m!n!}{i!(s!)^i}(-1){j-i}\frac{1}{(j-i)!(m-j)!(n-js)!}\frac{1}{(s!)^{j-i}}(m-j){n-js}$

再合并一下各个分数线

$ans=\sum_{j=0}^{{lim}\sum_{i=0}}{j}(-1)^{{j-i}\frac{w_im!n!}{i!(s!)}j(j-i)!(m-j)!(n-js)!}(m-j)^{n-js}$

把所有只含$j$的提到前面一个sigma去

$ans=\sum_{j=0}^{{lim}\frac{m!n!(m-j)}{n-js}}{(s!)^{j(m-j)!(n-js)!}\sum_{i=0}}{j}\frac{(-1)^{j-i}w_i}{i!(j-i)!}$

发现后面是一个卷积的形式！

设$a(i)=\frac{w_i}{i!},b(i)=\frac{(-1)^i}{i!}$

那么：

$ans=\sum_{j=0}^{{lim}\frac{m!n!(m-j)}{n-js}}{(s!)^j(m-j)!(n-js)!}(a\ast b)(j)$

所有的逆元、阶乘都可以$O(n)$预处理，卷积使用模数为998244353的$NTT$来实现，总时间复杂度为$O(n+limlog_2lim)$

Code:

我的代码好像常数比较大......如果有知道怎么优化的dalao，麻烦在评论区帮忙指出一下！

不胜感激！

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define ll long long
using namespace std;
inline ll read(){
	ll re=0,flag=1;char ch=getchar();
	while(ch>'9'||ch<'0'){
		if(ch=='-') flag=-1;
		ch=getchar();
	}
	while(ch>='0'&&ch<='9') re=(re<<1)+(re<<3)+ch-'0',ch=getchar();
	return re*flag;
}
ll MOD=1004535809,g=3,f[10000010],finv[10000010],w[10000010],invlimit;
ll qpow(ll a,ll b){
	ll re=1;
	while(b){
		if(b&1) re=re*a%MOD;
		a=a*a%MOD;b>>=1;
	}
	return re;
}
ll A[400010],B[400010],r[400010],n,limit,cnt,m,S,N,invg;
void init(){
	limit=1;cnt=0;ll i;
	while(limit<=(N<<1)) limit<<=1,cnt++;
	for(i=0;i<limit;i++) r[i]=((r[i>>1]>>1)|((i&1)<<(cnt-1)));
	f[1]=finv[1]=f[0]=finv[0]=1;
	for(i=2;i<=max(max(n,m),S);i++){
		f[i]=f[i-1]*i%MOD;
	}
	finv[max(max(n,m),S)]=qpow(f[max(max(n,m),S)],MOD-2);
	for(i=max(max(n,m),S);i>=1;i--) finv[i-1]=finv[i]*i%MOD;
	invlimit=qpow(limit,MOD-2);invg=qpow(g,MOD-2);
}
void ntt(ll *a,ll type){
	ll i,j,k,mid,len;ll y,w,wn;
	for(i=0;i<limit;i++) if(i<r[i]) swap(a[i],a[r[i]]);
	for(mid=1;mid<limit;mid<<=1){
		len=mid<<1;
		wn=qpow(((type==1)?g:invg),(MOD-1)/len);
		for(j=0;j<limit;j+=len){
			w=1;
			for(k=0;k<mid;k++,w=w*wn%MOD){
				y=a[j+k+mid]*w%MOD;
				a[j+k+mid]=(a[j+k]-y+MOD)%MOD;
				a[j+k]=(a[j+k]+y)%MOD;
			}
		}
	}
	if(type==-1) for(i=0;i<limit;i++) a[i]=(a[i]*invlimit%MOD);
}
int main(){
	n=read();m=read();S=read();N=min(n/S,m);ll i;
	for(i=0;i<=m;i++) w[i]=read();
	init();
	for(i=0;i<=N;i++){
		A[i]=finv[i]*w[i]%MOD;
		B[i]=((i&1)?MOD-1:1)*finv[i]%MOD;
	}
	ntt(A,1);ntt(B,1);
	for(i=0;i<limit;i++) A[i]=A[i]*B[i]%MOD;
	ntt(A,-1);
	ll ans=0;
	for(i=0;i<=N;i++){
		ans=(ans+(finv[m-i]*finv[n-i*S]%MOD*qpow(finv[S],i)%MOD*qpow(m-i,n-i*S)%MOD*A[i]))%MOD;
	}
	printf("%lld\n",ans*f[n]%MOD*f[m]%MOD);
}