洛谷 P4491 [HAOI2018]染色 ntt

最新推荐文章于 2023-10-15 16:06:08 发布

Amber_lylovely

最新推荐文章于 2023-10-15 16:06:08 发布

阅读量207

点赞数

分类专栏： fft-ntt

本文链接：https://blog.csdn.net/liangzihao1/article/details/88722873

版权

fft-ntt 专栏收录该内容

27 篇文章 0 订阅

订阅专栏

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

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

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

输入输出格式

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

接下来一行 $M + 1$ 个整数, 第 $i$ 个数表示 $W_{i-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$
对于 $100\%$ 的数据, 满足 $\le W_i < 1004535809$

Data
在这里插入图片描述
分析：
设 $f [i]$ 表示至少有 $i$ 种颜色相同的方案数。
首先要在 $m$ 中颜色中选出 $i$ 中，方案为 $\binom{m}{i}$ 。对于这 $i$ 种颜色，考虑每一种的放到画布上的方案，显然第一种是 $\binom{n}{s}$ ，第二种是 $\binom{n-s}{s}$ ，……，乘起来就是 $\frac{n!}{(s!)^i*(n-s*i)!}$ ，最后剩余的直接都是可以乱选的。
所以，
$f[i]=\binom{m}{i}*\frac{n!}{(s!)^i*(n-s*i)!}*(m-i)^{(n-s*i)}$
设恰好为 $i$ 的方案为 $g [i]$ ，直接容斥得到，
$g[i]=\sum_{j=i}^{m}(-1)^{j-i}*\binom{j}{i}*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);
}