【BZOJ5306】染色（HAOI2018）-容斥原理+NTT

测试地址：染色
做法：本题需要用到容斥原理+NTT。
好吧，我承认以下的推导过程是借（chao）鉴（xi）这位大佬的，Orz。
要求恰有i(0≤i≤E,E=min(⌊nS⌋,m))i(0≤i≤E,E=min(⌊nS⌋,m))种颜色出现SS次的方案数，其实就是要求其他m−im−i种颜色必定不能恰好出现SS次，用容斥原理列出式子得：
ans=∑Ei=0WiCimCiSn(iS)!(S!)i∑E−ij=0(−1)jCjm−iCjSn−iS(jS)!(S!)j(m−i−j)n−iS−jSans=∑i=0EWiCmiCniS(iS)!(S!)i∑j=0E−i(−1)jCm−ijCn−iSjS(jS)!(S!)j(m−i−j)n−iS−jS
在第二个和式中用j−ij−i替换jj，则有：
ans=∑Ei=0WiCimCiSn(iS)!(S!)i∑Ej=i(−1)j−iCj−im−iC(j−i)Sn−iS((j−i)S)!(S!)j−i(m−j)n−jSans=∑i=0EWiCmiCniS(iS)!(S!)i∑j=iE(−1)j−iCm−ij−iCn−iS(j−i)S((j−i)S)!(S!)j−i(m−j)n−jS
把组合数拆开，化简得：
ans=∑Ei=0Wimni!∑Ej=i(−1)j−i(m−j)n−jS(j−i)!(m−j)!(n−jS)!(S!)jans=∑i=0EWimni!∑j=iE(−1)j−i(m−j)n−jS(j−i)!(m−j)!(n−jS)!(S!)j
交换i,ji,j的位置，整理得：
ans=∑Ej=0m!n!(m−j)n−jS(m−j)!(n−jS)!(S!)j∑ji=0Wii!×(−1)j−i(j−i)!ans=∑j=0Em!n!(m−j)n−jS(m−j)!(n−jS)!(S!)j∑i=0jWii!×(−1)j−i(j−i)!
显然后半部分是一个卷积的形式，使用NTT求出即可。
以下是本人代码：

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll mod=1004535809;
const ll g=3;
int r[400010];
ll n,m,s,E,fac[10000010],inv[10000010];
ll a[400010],b[400010],w[100010];

ll power(ll a,ll b)
{
    ll s=1,ss=a;
    while(b)
    {
        if (b&1) s=s*ss%mod;
        ss=ss*ss%mod;b>>=1;
    }
    return s;
}

void NTT(ll *a,ll type,int n)
{
    for(int i=0;i<n;i++)
        if (i<r[i]) swap(a[i],a[r[i]]);
    for(int mid=1;mid<n;mid<<=1)
    {
        ll W=power(g,(type*(mod-1)/(mid<<1)%(mod-1)+(mod-1))%(mod-1));
        for(int l=0;l<n;l+=(mid<<1))
        {
            ll w=1;
            for(int k=0;k<mid;k++,w=w*W%mod)
            {
                ll x=a[l+k],y=w*a[l+mid+k]%mod;
                a[l+k]=(x+y)%mod;
                a[l+mid+k]=(x-y+mod)%mod;
            }
        }
    }
    if (type==-1)
    {
        ll inv=power(n,mod-2);
        for(int i=0;i<n;i++)
            a[i]=a[i]*inv%mod;
    }
}

int main()
{
    scanf("%lld%lld%lld",&n,&m,&s);
    for(int i=0;i<=m;i++)
        scanf("%lld",&w[i]);
    fac[0]=fac[1]=inv[0]=inv[1]=1;
    for(ll i=2;i<=max(n,m);i++)
    {
        fac[i]=fac[i-1]*i%mod;
        inv[i]=(mod-mod/i)*inv[mod%i]%mod;
    }
    for(ll i=2;i<=max(n,m);i++)
        inv[i]=inv[i]*inv[i-1]%mod;

    E=min(n/s,m);
    for(ll i=0;i<=E;i++)
    {
        a[i]=w[i]*inv[i]%mod;
        b[i]=(((i%2)?-1:1)*inv[i]%mod+mod)%mod;
    }

    int x=1,bit=0;
    while(x<=(E<<1)) x<<=1,bit++;
    r[0]=0;
    for(int i=1;i<=x;i++)
        r[i]=(r[i>>1]>>1)|((i&1)<<(bit-1));
    NTT(a,1ll,x),NTT(b,1ll,x);
    for(int i=0;i<=x;i++)
        a[i]=a[i]*b[i]%mod;
    NTT(a,-1ll,x);

    ll ans=0;
    for(ll j=0;j<=E;j++)
    {
        ll now=1;
        now=fac[m]*fac[n]%mod*inv[m-j]%mod*inv[n-j*s]%mod;
        now=now*power(m-j,n-j*s)%mod;
        now=now*power(power(fac[s],j),mod-2)%mod;
        now=now*a[j]%mod;
        ans=(ans+now)%mod;
    }
    printf("%lld",ans);

    return 0;
}