洛谷P4491：[HAOI2018]染色（容斥+ntt）

今年 HAOI 好强
题面

设$H=min(\frac{n}{s},m)$
从0到H 枚举题意中的k
再枚举哪k种颜色，放哪里

然后看题解

剩下的就是m-k种颜色，n-sk个位置，恰好0种颜色出现恰好s次的方案数
额，容斥把恰好转为至少
就是枚举至少j种，哪j种，放哪里，剩下的随便放

借(dao)鉴(yong)别人的柿子

$$ans=\sum_{i=0}^Nw[i]{m\choose i}{n\choose is}\frac{(is)!}{(s!)^i}\sum_{j=0}^{N-i}(-1)^j{m-i\choose j}{n-is\choose js}\frac{(js)!}{(s!)^j}(m-i-j)^{n-is-js}$$

看模数是998244353以外的费马素数
大概是ntt的题
套路拆组合数
把$j$用$j+i$换，拉到前面枚举

$$\sum_{j=0}^{N}\frac{m!n!}{(m-j)!(n-js)!}(\frac{1}{s!})^j(m-j)^{n-js}\sum_{i=0}^j\frac{w[i]}{i!}*\frac{(-1)^{j-i}}{(j-i)!}$$
后面就是个卷积

#include <iostream>
#include <fstream>
#include <algorithm>
#include <cmath>
#include <ctime>
#include <cstdio>
#include <cstdlib>
#include <cstring>

using namespace std;
#define mmst(a, b) memset(a, b, sizeof(a))
#define mmcp(a, b) memcpy(a, b, sizeof(b))

typedef long long LL;

const int N=1000500,BN=10000010;
const LL p=1004535809;

int n,rev[N];
LL nn,njc=1,H,m,s;
LL jc[BN],I[BN],Ijc[BN];
LL w[N];
LL aa[N],bb[N];
LL ans;

LL cheng(LL a,LL b)
{
    LL res=1;
    for(;b;b>>=1,a=a*a%p)
    if(b&1)
    res=res*a%p;
    return res;
}

void init(int lim)
{
    int k=-1;
    n=1;
    while(n<=lim)
    k++,n<<=1;

    for(int i=0;i<n;i++)
    rev[i]=(rev[i>>1] >> 1) | ((i&1)<<k);
}

void ntt(LL *a,int ops)
{
    for(int i=0;i<n;i++)
    if(i<rev[i])
    swap(a[i],a[rev[i]]);

    for(int m=1,l=2;m<n;m<<=1,l<<=1)
    {
        LL wn= (ops) ? cheng(3,(p-1)/l) : cheng(3,p-1-(p-1)/l);
        for(int i=0;i<n;i+=l)
        {
            LL w=1;
            for(int k=0;k<m;k++)
            {
                LL t=a[i+k+m]*w%p;
                a[i+k+m]=(a[i+k]-t+p)%p;
                a[i+k]=(a[i+k]+t)%p;
                w=w*wn%p;
            }
        }
    }

    if(!ops)
    for(int i=0;i<n;i++)
    a[i]=a[i]*I[n]%p;
}

int main()
{
    cin>>nn>>m>>s;

    if(s==0)
    H=m;
    else
    H=min(nn/s,m);

    for(int i=0;i<=m;i++)
    scanf("%lld",&w[i]);

    I[1]=jc[0]=Ijc[0]=1;

    for(int i=2;i<BN;i++)
    I[i]=I[p%i]*(p-p/i)%p;

    for(int i=1;i<BN;i++)
    jc[i]=jc[i-1]*i%p,Ijc[i]=Ijc[i-1]*I[i]%p;

    for(int i=0;i<=H;i++)
    {
        aa[i]=w[i]*Ijc[i]%p;
        if(i&1)
        bb[i]=(p-Ijc[i])%p;
        else
        bb[i]=Ijc[i];
    }

    init(H+H+5);
    ntt(aa,1);
    ntt(bb,1);
    for(int i=0;i<n;i++)
    aa[i]=aa[i]*bb[i]%p;
    ntt(aa,0);

    for(int j=0;j<=H;j++)
    ans=(ans+Ijc[m-j]*Ijc[nn-j*s]%p*cheng(Ijc[s],j)%p*cheng(m-j,nn-j*s)%p*aa[j])%p;

    ans=ans*jc[nn]%p*jc[m]%p;

    cout<<ans<<endl;

    return 0;
}