HAOI 2018 染色（容斥+NTT）

题意

https://loj.ac/problem/2527

思路

设 (f(k)) 为强制选择 (k) 个颜色出现 (s) 种，其余任取的方案数。

则有

[f(k)={mchoose k}{nchoose sk}{(sk)!over(s!)^k}(m-k)^{n-sk} ]

不难看出，这个方案可能包括了超过 (k) 种颜色，也有重复的方案，所以恰有 (k) 个颜色出现 (s) 种的方案 (ans_k) 满足

[ans_k=sum_{i=k}^{min(m,{nover s})}(-1)^{i-k}{ichoose k}f(i) ]

最终化简得到

[ans_k={1over k!}sum_{i=k}^{min(m,{nover s})}i!f(i)cdot {(-1)^{i-k}over (i-k)!} ]

( ext{NTT}) 卷积即可。

做容斥题时定义出的状态本身就是有重复的，所以需要加加减减。

代码

#include<bits/stdc++.h>
#define FOR(i,x,y) for(int i=(x),i##END=(y);i<=i##END;++i)
#define DOR(i,x,y) for(int i=(x),i##END=(y);i>=i##END;--i)
using namespace std;
template<typename T,typename _T>inline bool chk_min(T &x,const _T y){return y<x?x=y,1:0;}
template<typename T,typename _T>inline bool chk_max(T &x,const _T y){return x<y?x=y,1:0;}
typedef long long ll;
const int P=1004535809;
const int N=1e7+5;
const int M=1<<17|5;
namespace _Maths
{
	ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
	void exgcd(ll a,ll b,ll &x,ll &y)
	{
		if(!b){x=1,y=0;return;}
		exgcd(b,a%b,y,x);y-=a/b*x;
	}
	ll Pow(ll a,ll p,ll P)
	{
		ll res=1;
		for(;p>0;p>>=1,(a*=a)%=P)if(p&1)(res*=a)%=P;
		return res;
	}
	ll inv(ll a,ll P){ll x,y;exgcd(a,P,x,y);return (x%P+P)%P;}
}
using namespace _Maths;
namespace _Polynomial
{
	int A[M<<1],B[M<<1],C[M<<1];
	int w[M<<1],r[M<<1];
	void DFT(int *a,int op,int n)
	{
		FOR(i,0,n-1)if(i<r[i])swap(a[i],a[r[i]]);
		for(int i=2;i<=n;i<<=1)
			for(int j=0;j<n;j+=i)
				for(int k=0;k<i/2;k++)
				{
					int u=a[j+k],t=(ll)w[op==1?n/i*k:(n-n/i*k)&(n-1)]*a[j+k+i/2]%P;
					a[j+k]=(u+t)%P,a[j+k+i/2]=(u-t)%P;
				}
		if(op==-1)
		{
			int I=inv(n,P);
			FOR(i,0,n-1)a[i]=(ll)a[i]*I%P;
		}
	}
	void multiply(const int *a,const int *b,int *c,int n1,int n2)
	{
		int n=1;
		while(n<n1+n2-1)n<<=1;
		FOR(i,0,n1-1)A[i]=a[i];
		FOR(i,0,n2-1)B[i]=b[i];
		FOR(i,n1,n-1)A[i]=0;
		FOR(i,n2,n-1)B[i]=0;
		FOR(i,0,n-1)r[i]=(r[i>>1]>>1)|((i&1)*(n>>1));
		w[0]=1,w[1]=Pow(3,(P-1)/n,P);
		FOR(i,2,n-1)w[i]=(ll)w[i-1]*w[1]%P;
		
		DFT(A,1,n),DFT(B,1,n);
		FOR(i,0,n-1)A[i]=(ll)A[i]*B[i]%P;
		DFT(A,-1,n);
		FOR(i,0,n1+n2-2)c[i]=(A[i]+P)%P;
	}
};
int fac[N],ifac[N],f[M];
int A[M],B[M],C[M<<1];
int W[M];
int n,m,s,b;
ll ans;

ll Com(int n,int m){return (ll)fac[n]*ifac[m]%P*ifac[n-m]%P;}

int main()
{
	fac[0]=fac[1]=1;FOR(i,2,N-1)fac[i]=(ll)fac[i-1]*i%P;
	ifac[0]=ifac[1]=1;FOR(i,2,N-1)ifac[i]=(ll)(P-P/i)*ifac[P%i]%P;
	FOR(i,2,N-1)ifac[i]=(ll)ifac[i-1]*ifac[i]%P;
	scanf("%d%d%d",&n,&m,&s);
	FOR(i,0,m)scanf("%d",&W[i]);
	b=min(m,n/s);
	FOR(i,0,b)f[i]=Com(m,i)*Com(n,s*i)%P*fac[s*i]%P*Pow(ifac[s],i,P)%P*Pow(m-i,n-s*i,P)%P;
	FOR(i,0,b)A[i]=(ll)fac[i]*f[i]%P;
	FOR(i,-b,0)B[i+b]=Pow(-1,-i,P)*ifac[-i]%P;
	_Polynomial::multiply(A,B,C,b+1,b+1);
	FOR(i,0,b)(ans+=(ll)W[i]*ifac[i]%P*C[i+b]%P)%=P;
	printf("%lld
",(ans%P+P)%P);
	return 0;
}