LibreOJ #2527.「HAOI2018」染色 容斥原理+FFT

题意

为了报答小C的苹果, 小G打算送给热爱美术的小C一块画布, 这块画布可以抽象为一个长度为$n$的序列, 每个位置都可以被染成$m$种颜色中的某一种.
然而小C只关心序列的$n$个位置中出现次数恰好为$s$的颜色种数, 如果恰好出现了$s$次的颜色有$k$种, 则小C会产生$w_k$的愉悦度.
小C希望知道对于所有可能的染色方案, 他能获得的愉悦度的和对$1004535809$取模的结果是多少.
$n\le10^7,m\le10^5$

分析

先考虑对于每个$k$如何求出有多少种染色方案满足恰好有$k$种颜色出现次数为$S$，设为$ans_k$。
设$f(k)$表示至少有$k$种颜色出现次数恰好为$S$的方案数，显然

$$f(k)=C_m^kC_n^{ks}\frac{(ks)!}{(s!)^k}(m-k)^{n-ks}$$
根据容斥原理不难得到 $$ans_k=\sum_{i=k}^m(-1)^{i-k}C_i^kf(i)$$
展开之后发现是卷积形式，用NTT优化即可。

代码

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>

typedef long long LL;

const int N=400005;
const int MOD=1004535809;

int n,m,s,jc[10000005],ny[10000005],f[N],g[N],rev[N],L;

int read()
{
    int x=0,f=1;char ch=getchar();
    while (ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while (ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}

int ksm(int x,int y)
{
    int ans=1;
    while (y)
    {
        if (y&1) ans=(LL)ans*x%MOD;
        x=(LL)x*x%MOD;y>>=1;
    }
    return ans;
}

int C(int n,int m)
{
    return (LL)jc[n]*ny[m]%MOD*ny[n-m]%MOD;
}

void NTT(int *a,int f)
{
    for (int i=0;i<L;i++) if (i<rev[i]) std::swap(a[i],a[rev[i]]);
    for (int i=1;i<L;i<<=1)
    {
        int wn=ksm(3,f==1?(MOD-1)/i/2:MOD-1-(MOD-1)/i/2);
        for (int j=0;j<L;j+=(i<<1))
        {
            int w=1;
            for (int k=0;k<i;k++)
            {
                int u=a[j+k],v=(LL)a[j+k+i]*w%MOD;
                a[j+k]=(u+v)%MOD;a[j+k+i]=(u+MOD-v)%MOD;
                w=(LL)w*wn%MOD;
            }
        }
    }
    int ny=ksm(L,MOD-2);
    if (f==-1) for (int i=0;i<L;i++) a[i]=(LL)a[i]*ny%MOD;
}

int main()
{
    n=read();m=read();s=read();
    jc[0]=jc[1]=ny[0]=ny[1]=1;
    for (int i=2;i<=std::max(n,m);i++) jc[i]=(LL)jc[i-1]*i%MOD,ny[i]=(LL)(MOD-MOD/i)*ny[MOD%i]%MOD;
    for (int i=2;i<=std::max(n,m);i++) ny[i]=(LL)ny[i-1]*ny[i]%MOD;
    for (int k=0;k<=m&&k*s<=n;k++) f[k]=(LL)C(m,k)*C(n,k*s)%MOD*jc[k*s]%MOD*ksm(ny[s],k)%MOD*ksm(m-k,n-k*s)%MOD*jc[k]%MOD;
    for (int i=0;i<=m;i++)
        if ((m-i)&1) g[i]=MOD-ny[m-i];
        else g[i]=ny[m-i];
    int lg=0;
    for (L=1;L<=m*2;L<<=1,lg++);
    for (int i=0;i<L;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(lg-1));
    NTT(f,1);NTT(g,1);
    for (int i=0;i<L;i++) f[i]=(LL)f[i]*g[i]%MOD;
    NTT(f,-1);
    int ans=0;
    for (int i=0;i<=m;i++)
    {
        int x=read();
        (ans+=(LL)f[m+i]*x%MOD*ny[i]%MOD)%=MOD;
    }
    printf("%d",ans);
    return 0;
}