[BZOJ5306][HAOI2018]染色(容斥+FFT)

https://www.cnblogs.com/zhoushuyu/p/9138251.html

注意如果一开始F(i)中内层式子中j枚举的是除前i种颜色之外还有几种出现S次的颜色，那么后面式子就会难推很多。

 #include<cstdio>
 #include<algorithm>
 #define rep(i,l,r) for (int i=(l); i<=(r); i++)
 using namespace std;

 const int N=,M=,mod=;
 int n,m,s,ans,w[N],fac[M],inv[M],rev[N],a[N],b[N];

 int ksm(int a,int b){
     int res=;
     for (; b; a=1ll*a*a%mod,b>>=)
         if (b & ) res=1ll*res*a%mod;
     return res;
 }

 void NTT(int a[],int n,bool f){
     for (int i=; i<n; i++) if (i<rev[i]) swap(a[i],a[rev[i]]);
     for (int i=; i<n; i<<=){
         int wn=ksm(,f ? (mod-)/(i<<) : (mod-)-(mod-)/(i<<));
         for (int p=i<<,j=; j<n; j+=p){
             int w=;
             for (int k=; k<i; k++,w=1ll*w*wn%mod){
                 int x=a[j+k],y=1ll*w*a[i+j+k]%mod;
                 a[j+k]=(x+y)%mod; a[i+j+k]=(x-y+mod)%mod;
             }
         }
     }
     if (f) return;
     int inv=ksm(n,mod-);
     for (int i=; i<n; i++) a[i]=1ll*a[i]*inv%mod;
 }

 int main(){
     freopen("color.in","r",stdin);
     freopen("color.out","w",stdout);
     scanf("%d%d%d",&n,&m,&s); int N=min(m,n/s),ed=max(n,m);
     rep(i,,m) scanf("%d",&w[i]);
     fac[]=; rep(i,,ed) fac[i]=1ll*fac[i-]*i%mod;
     inv[ed]=ksm(fac[ed],mod-);
     for (int i=ed-; ~i; i--) inv[i]=1ll*inv[i+]*(i+)%mod;
     rep(i,,N) a[i]=1ll*w[i]*inv[i]%mod;
     rep(i,,N) b[i]=(i&)?mod-inv[i]:inv[i];
     int nn=,L=; for (; nn<=*N; nn<<=,L++);
     for (int i=; i<nn; i++) rev[i]=(rev[i>>]>>)|((i&)<<(L-));
     NTT(a,nn,); NTT(b,nn,);
     for (int i=; i<nn; i++) a[i]=1ll*a[i]*b[i]%mod;
     NTT(a,nn,);
     rep(i,,N) ans=(ans+1ll*ksm(m-i,n-i*s)*inv[m-i]%mod*ksm(inv[s],i)%mod*inv[n-i*s]%mod*a[i]%mod)%mod;
     printf("%lld\n",1ll*ans*fac[n]%mod*fac[m]%mod);
     return ;
 }