洛谷 P4491 [HAOI2018]染色 ntt

题目描述
为了报答小 C 的苹果, 小 G 打算送给热爱美术的小 C 一块画布, 这块画布可 以抽象为一个长度为 $N$N 的序列, 每个位置都可以被染成 $M$M 种颜色中的某一种。

然而小 C 只关心序列的 $N$N 个位置中出现次数恰好为 $S$S 的颜色种数, 如果恰 好出现了 $S$S 次的颜色有 $K$K 种, 则小 $C$C 会产生 $W_k$Wk​ 的愉悦度.

小 $C$C 希望知道对于所有可能的染色方案, 他能获得的愉悦度的和对 $1004535809$1004535809 取模的结果是多少.

输入输出格式

输入格式：
从标准输入读入数据. 第一行三个整数 $N, M, S$N,M,S

接下来一行 $M + 1$M+1 个整数, 第 $i$i 个数表示 $W_{i-1}$Wi−1​​ .

输出格式：
输出到标准输出中. 输出一个整数表示答案．
输入输出样例

输入样例#1：
8 8 3
3999 8477 9694 8454 3308 8961 3018 2255 4910
输出样例#1：
524070430

说明

特殊性质: $\forall 1 \le i \le m, W_i = 0$∀1≤i≤m,Wi​=0
对于 $100\%$100% 的数据, 满足 $0 \le W_i &lt; 1004535809$0≤Wi​<1004535809

Data

分析：
设$f[i]$f[i]表示至少有$i$i种颜色相同的方案数。
首先要在$m$m中颜色中选出$i$i中，方案为$\binom{m}{i}$(im​)。对于这$i$i种颜色，考虑每一种的放到画布上的方案，显然第一种是$\binom{n}{s}$(sn​)，第二种是$\binom{n-s}{s}$(sn−s​)，……，乘起来就是$\frac{n!}{(s!)^i*(n-s*i)!}$(s!)i∗(n−s∗i)!n!​，最后剩余的直接都是可以乱选的。
所以，
$f[i]=\binom{m}{i}*\frac{n!}{(s!)^i*(n-s*i)!}*(m-i)^{(n-s*i)}$f[i]=(im​)∗(s!)i∗(n−s∗i)!n!​∗(m−i)(n−s∗i)
设恰好为$i$i的方案为$g[i]$g[i]，直接容斥得到，
$g[i]=\sum_{j=i}^{m}(-1)^{j-i}*\binom{j}{i}*f[j]$g[i]=j=i∑m​(−1)j−i∗(ij​)∗f[j]
把组合数拆开然后NTT即可。

代码：

// luogu-judger-enable-o2
#include <iostream>
#include <cstdio>
#include <cmath>
#include <algorithm>
#define LL long long

const int N=1e7+7;
const int maxn=3e5+7;
const LL G=3;
const LL mod=1004535809;

using namespace std;

int n,m,S,len;
LL a[maxn],jc[N],inv[N],f[maxn],g[maxn],r[maxn],w[maxn];

LL ksm(LL x,LL y)
{
    if (y==0) return 1;
    LL c=ksm(x,y/2);
    c=(c*c)%mod;
    if (y&1) c=(c*x)%mod;
    return c;
}

void prework()
{
    jc[0]=1;
    for (int i=1;i<=max(n,m);i++) jc[i]=jc[i-1]*(LL)i%mod;
    inv[max(n,m)]=ksm(jc[max(n,m)],mod-2);
    for (int i=max(n,m);i>0;i--) inv[i-1]=inv[i]*(LL)i%mod;
}

void ntt(LL *a,int f)
{
    for (int i=0;i<len;i++)
    {
        if (i<r[i]) swap(a[i],a[r[i]]);
    }
    w[0]=1;
    for (int i=2;i<=len;i<<=1)
    {
        LL wn;
        if (f==1) wn=ksm(G,(mod-1)/i);
             else wn=ksm(G,(mod-1)-(mod-1)/i);
        for (int j=i/2-2;j>=0;j-=2) w[j]=w[j/2];
        for (int j=1;j<i/2;j+=2) w[j]=(w[j-1]*wn)%mod;
        for (int j=0;j<len;j+=i)
        {
            for (int k=0;k<i/2;k++)
            {
                LL u=a[j+k],v=a[j+k+i/2]*w[k]%mod;
                a[j+k]=(u+v)%mod;
                a[j+k+i/2]=(u+mod-v)%mod;
            }
        }
    }
    if (f==-1)
    {
        LL inv=ksm(len,mod-2);
        for (int i=0;i<len;i++) a[i]=a[i]*inv%mod;
    }
}

void NTT(LL *a,LL *b,LL *c,int n,int m)
{
    len=1;
    int k=0;
    while (len<=n+m) len<<=1,k++;
    for (int i=0;i<len;i++) r[i]=(r[i>>1]>>1)|((i&1)<<(k-1));
    ntt(a,1),ntt(b,1);
    for (int i=0;i<len;i++) c[i]=a[i]*b[i]%mod;
    ntt(c,-1);
}

int main()
{
    scanf("%d%d%d",&n,&m,&S);
    for (int i=0;i<=m;i++) scanf("%lld",&a[i]);
    prework();	
    for (int i=0;i<=m;i++)
    {
        if (i*S>n) break;
        f[i]=jc[m]*inv[m-i]%mod*jc[n]%mod*ksm(inv[S],i)%mod*inv[n-S*i]%mod*ksm(m-i,n-S*i)%mod;
    }
    for (int i=0;i<=m;i++) 
    {
        if (i&1) g[i]=mod-inv[i];
            else g[i]=inv[i];
    }
    reverse(f,f+m+1);
    NTT(f,g,f,m+1,m+1);
    reverse(f,f+m+1);
    LL ans=0;
    for (int i=0;i<=m;i++) ans=(ans+f[i]*a[i]%mod*inv[i]%mod)%mod;
    printf("%lld\n",ans);
}