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

测试地址：染色
做法：本题需要用到容斥原理+NTT。
好吧，我承认以下的推导过程是借（chao）鉴（xi）这位大佬的，Orz。
要求恰有i(0≤i≤E,E=min(⌊nS⌋,m))$i(0\le i\le E,E=min(\lfloor \frac{n}{S}\rfloor ,m))$种颜色出现S$S$次的方案数，其实就是要求其他m−i$m-i$种颜色必定不能恰好出现S$S$次，用容斥原理列出式子得：
ans=∑Ei=0WiCimCiSn(iS)!(S!)i∑E−ij=0(−1)jCjm−iCjSn−iS(jS)!(S!)j(m−i−j)n−iS−jS$ans=\sum _{i=0}^E{W}_i{C}_m^i{C}_n^{iS}\frac{(iS)!}{(S!{)}^i}\sum _{j=0}^{E-i}(-1{)}^j{C}_{m-i}^j{C}_{n-iS}^{jS}\frac{(jS)!}{(S!{)}^j}(m-i-j{)}^{n-iS-jS}$
在第二个和式中用j−i$j-i$替换j$j$，则有：
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−jS$ans=\sum _{i=0}^E{W}_i{C}_m^i{C}_n^{iS}\frac{(iS)!}{(S!{)}^i}\sum _{j=i}^E(-1{)}^{j-i}{C}_{m-i}^{j-i}{C}_{n-iS}^{(j-i)S}\frac{((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!)j$ans=\sum _{i=0}^E{W}_i\frac{mn}{i!}\sum _{j=i}^E\frac{(-1{)}^{j-i}(m-j{)}^{n-jS}}{(j-i)!(m-j)!(n-jS)!(S!{)}^j}$
交换i,j$i,j$的位置，整理得：
ans=∑Ej=0m!n!(m−j)n−jS(m−j)!(n−jS)!(S!)j∑ji=0Wii!×(−1)j−i(j−i)!$ans=\sum _{j=0}^E\frac{m!n!(m-j{)}^{n-jS}}{(m-j)!(n-jS)!(S!{)}^j}\sum _{i=0}^j\frac{W_i}{i!}\times \frac{(-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;
}