洛谷P4491 [HAOI2018]染色

题面
题解：首先这个“恰好”看着很不爽，换成“至少”。
设\(f[i]\)表示颜色个数为\(S\)的颜色至少有\(i\)个的方案数。
考虑如何计算。
1.\(m\)个颜色选了\(i\)个，\(\binom{m}{i}\)；
2.\(i\)个颜色选了恰好\(S\)个，其他的\(m-i\)个颜色任选;$$\frac {n!}{ {S!}^i \times fac[n-i*S]} $$
3.剩下\(n-i*S\)个位置任选\(m-i\)中颜色，\({m-i}^{n-i*S}\)。
把它们乘起来就是\(f[i]\)。
接下来考虑计算答案。
考虑容斥。设\(g[i]\)为颜色个数为\(S\)的颜色恰好有\(i\)个的方案数。于是有：

\[g[i]= \sum_{j=i}^{lim} \frac { {-1}^{j-i} } {fac[j-i]} \times f[j] \times \binom{j}{i} \]

这个\(lim\)表示的是最多能取到的恰好有\(S\)个的颜色数，\(lim=min(n/S,m)\)。
将组合数拆开，试着把式子化为卷积的形式。

\[g[i] \times fac[i] = \sum_{j=i}^{lim} \frac { {-1}^{j-i} } {fac[j-i]} \times f[j] \times fac[j] \]

这样用NTT做就好了。
令$$a[i]=f[i] \times fac[i],b[i]= \frac { {-1}^{j-i} } {fac[j-i]} $$
把\(b\)数组翻转一下，取做完卷积的数组的\(lim\)到\(2lim\)位更新答案即可。
时间复杂度：\(O(n+mlogm)\)
代码：

#include<bits/stdc++.h>
using namespace std;
#define re register int
#define F(x,y,z) for(re x=y;x<=z;x++)
#define FOR(x,y,z) for(re x=y;x>=z;x--)
typedef long long ll;
#define I inline void
#define IN inline int
#define Cl(x,y) memset(x,y,sizeof(x))
#define kTk system("pause")
template<class D>I read(D &res){
	res=0;register D g=1;register char ch=getchar();
	while(!isdigit(ch)){
		if(ch=='-')g=-1;
		ch=getchar();
	}
	while(isdigit(ch)){
		res=(res<<3)+(res<<1)+(ch^48);
		ch=getchar();
	}
	res*=g;
}
const int Mod=1004535809;
int n,m,k,S,len,lim,ans,r[303000],w[303000],fac[10100000],inv[10100000],f[303000],A[303000],B[303000];
I Add(int &x,int y){
	x+=y;if(x>=Mod)x-=Mod;
}
IN Plus(int x,int y){
	x+=y;if(x>=Mod)x-=Mod;return x;
}
IN Pow(int x,int y){
	re res=1;
	while(y){
		if(y&1)res=(ll)res*x%Mod;
		x=(ll)x*x%Mod;
		y>>=1;
	}
	return res;
}
IN C(int x,int y){
	return (ll)fac[x]*inv[y]%Mod*inv[x-y]%Mod;
}
I init(){
	re M=10001000;
	fac[0]=1;
	F(i,1,M)fac[i]=(ll)fac[i-1]*i%Mod;
	inv[M]=Pow(fac[M],Mod-2);
	FOR(i,M-1,0)inv[i]=(ll)inv[i+1]*(i+1)%Mod;
}
I ntt(int *a,int sn){
	F(i,1,S)if(i<r[i])swap(a[i],a[r[i]]);
	for(re i=1;i<S;i<<=1){
		re gn=Pow(3,(Mod-1)/(i<<1));if(sn==-1)gn=Pow(gn,Mod-2);
		for(re p=i<<1,j=0;j<S;j+=p){
			re g0=1;
			for(re k=0;k<i;k++,g0=(ll)g0*gn%Mod){
				re X=a[j+k],Y=(ll)a[i+j+k]*g0%Mod;
				a[j+k]=Plus(X,Y);a[i+j+k]=Plus(X,Mod-Y);
			}
		}
	}
	if(sn==-1){
		re invn=Pow(S,Mod-2);
		F(i,0,S-1)a[i]=(ll)a[i]*invn%Mod;
	}
}
int main(){
	read(n);read(m);read(k);lim=min(n/k,m);
	F(i,0,m)read(w[i]);
	S=1;len=0;
	while(S<((lim<<1)+2))S<<=1,len++;
	init();
	F(i,0,lim)f[i]=(ll)C(m,i)*fac[n]%Mod*Pow(inv[k],i)%Mod*inv[n-i*k]%Mod*Pow(m-i,n-i*k)%Mod;
	F(i,0,S-1){
		A[i]=(ll)f[i]*fac[i]%Mod;
		B[i]=inv[lim-i];if((lim-i)&1)B[i]=(Mod-B[i])%Mod;
		r[i]=(r[i>>1]>>1)|((i&1)<<(len-1));
	}
	ntt(A,1);ntt(B,1);
	F(i,0,S-1)A[i]=(ll)A[i]*B[i]%Mod;
	ntt(A,-1);
	F(i,lim,lim<<1)Add(ans,(ll)A[i]*w[i-lim]%Mod*inv[i-lim]%Mod);
	printf("%d",ans);
	return 0;
}