BZOJ 5306: [Haoi2018]染色 二项式反演+NTT

给定长度为 $n$ 的序列, 每个位置都可以被染成 $m$ 种颜色中的某一种. 如果恰好出现了 $s$ 次的颜色有 $k$ 种, 则会产生 $w_{k}$ 的价值. 求对于所有可能的染色方案，获得价值和对 $1004535809$ 取模的结果.

设 $lim=min(m,\frac{n}{s})$，即最大可能的颜色出现种类.
按照套路，令 $f[i]$ 表示钦定 $i$ 种长度为 $s$ 出现的方案数，$g[i]$ 表示恰好 $i$ 种出现的方案数.
$f[k]=\binom{m}{k}\frac{n!}{(n-ks)!(s!)^k}(m-k)^{n-ks}$
组合意义就是选 $ks$ 个位置放出现次数为 $s$ 的颜色，然后其余部分随便放.
而 $g[k]=\sum_{i=k}^{lim}\binom{i}{k}(-1)^{i-k}f[i]$
因为我们要算贡献，所以要求 $g[1]....g[lim]$ ，而上面的式子是 $O(lim^2)$ 的.
考虑优化：
将上面 $\binom{i}{k}$ 展开，得 $g[k]=\frac{1}{k!}\sum_{i=k}^{lim} \frac{(-1)^{i-k}}{(i-k)!}f[i]\times i!$
令 $a[i]=\frac{(-1)^i}{i!}$，$b[i]=f[i]\times i!$ ，则 $g[k]=\frac{1}{k!}\sum_{i=k}^{lim} a[i-k]\times b[i]$
这是一个标准的卷积形式！
直接用 NTT 加速即可.

code:

#include <bits/stdc++.h>
#define N 800005
#define LL long long
#define setIO(s) freopen(s".in","r",stdin)
using namespace std;
const LL G=3;
const LL mod=1004535809;
LL A[N],B[N],Ct[N],f[N],g[N],fac[10000008],inv[10000007],val[N];
LL qpow(LL x,LL y)
{
    LL tmp=1ll;
    for(;y;y>>=1,x=x*x%mod)
        if(y&1)
            tmp=tmp*x%mod;
    return tmp;
}
LL Inv(LL x) { return qpow(x,mod-2); }
void NTT(LL *a,int n,int flag)
{
    int i,j,k,mid;
    for(i=k=0;i<n;++i)
    {
        if(i>k) swap(a[i],a[k]);
        for(j=n>>1;(k^=j)<j;j>>=1);
    }
    for(mid=1;mid<n;mid<<=1)
    {
        LL wn=qpow(G,(mod-1)/(mid<<1));
        if(flag==-1) wn=Inv(wn);
        for(i=0;i<n;i+=(mid<<1))
        {
            LL w=1ll;
            for(j=0;j<mid;++j)
            {
                LL x=a[i+j],y=w*a[i+mid+j]%mod;
                a[i+j]=(x+y)%mod,a[i+j+mid]=(x-y+mod)%mod;
                w=w*wn%mod;
            }
        }
    }
    if(flag==-1)
    {
        LL re=Inv(n);
        for(i=0;i<n;++i) a[i]=a[i]*re%mod;
    }
}
LL C(int x,int y) { return fac[x]*inv[y]%mod*inv[x-y]%mod; }
int main()
{
    // setIO("input");
    int n,m,s,i,j,lim;
    scanf("%d%d%d",&n,&m,&s);
    for(i=0;i<=m;++i) scanf("%lld",&val[i]);
    lim=min(m,n/s);
    inv[0]=fac[0]=1ll;
    int pp=max(n,m);
    for(i=1;i<=pp;++i)
    {
        fac[i]=fac[i-1]*1ll*i%mod;
    }
    inv[max(n,m)]=qpow(fac[max(n,m)],mod-2);
    for(i=max(n,m)-1;i>=0;i--) inv[i]=inv[i+1]*(i+1)%mod;
    for(i=0;i<=lim;++i)
    {
        f[i]=C(m,i)*fac[n]%mod*inv[n-i*s]%mod*qpow(inv[s],i)%mod*qpow(m-i,n-i*s)%mod*fac[i]%mod;
    }
    for(i=0;i<=lim;++i) A[i]=(inv[i]*(i&1?-1:1)+mod)%mod;
    for(i=0;i<=lim;++i) B[lim-i]=f[i];
    LL ans=0ll;
    int tmp=1;
    while(tmp<=lim*2) tmp<<=1;
    NTT(A,tmp,1),NTT(B,tmp,1);
    for(i=0;i<tmp;++i) Ct[i]=A[i]*B[i]%mod;
    NTT(Ct,tmp,-1);
    for(i=0;i<=lim;++i) g[i]=Ct[lim-i]*inv[i]%mod;
    for(i=0;i<=lim;++i)
    {
        (ans+=g[i]*val[i]%mod)%=mod;
    }
    printf("%lld\n",ans);
    return 0;
}