HAOI 2018 染色（容斥+NTT）

Posted 2021-02-03 paulliant

题意

https://loj.ac/problem/2527

思路

设 (f(k)) 为强制选择 (k) 个颜色出现 (s) 种，其余任取的方案数。

则有
[ f(k)={mchoose k}{nchoose sk}{(sk)!over(s!)^k}(m-k)^{n-sk} ]
不难看出，这个方案可能包括了超过 (k) 种颜色，也有重复的方案，所以恰有 (k) 个颜色出现 (s) 种的方案 (ans_k) 满足
[ ans_k=sum_{i=k}^{min(m,{nover s})}(-1)^{i-k}{ichoose k}f(i) ]
最终化简得到
[ ans_k={1over k!}sum_{i=k}^{min(m,{nover s})}i!f(i)cdot {(-1)^{i-k}over (i-k)!} ]
( ext{NTT}) 卷积即可。

做容斥题时定义出的状态本身就是有重复的，所以需要加加减减。

代码

#include<bits/stdc++.h>
#define FOR(i,x,y) for(int i=(x),i##END=(y);i<=i##END;++i)
#define DOR(i,x,y) for(int i=(x),i##END=(y);i>=i##END;--i)
using namespace std;
template<typename T,typename _T>inline bool chk_min(T &x,const _T y){return y<x?x=y,1:0;}
template<typename T,typename _T>inline bool chk_max(T &x,const _T y){return x<y?x=y,1:0;}
typedef long long ll;
const int P=1004535809;
const int N=1e7+5;
const int M=1<<17|5;
namespace _Maths
{
    ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
    void exgcd(ll a,ll b,ll &x,ll &y)
    {
        if(!b){x=1,y=0;return;}
        exgcd(b,a%b,y,x);y-=a/b*x;
    }
    ll Pow(ll a,ll p,ll P)
    {
        ll res=1;
        for(;p>0;p>>=1,(a*=a)%=P)if(p&1)(res*=a)%=P;
        return res;
    }
    ll inv(ll a,ll P){ll x,y;exgcd(a,P,x,y);return (x%P+P)%P;}
}
using namespace _Maths;
namespace _Polynomial
{
    int A[M<<1],B[M<<1],C[M<<1];
    int w[M<<1],r[M<<1];
    void DFT(int *a,int op,int n)
    {
        FOR(i,0,n-1)if(i<r[i])swap(a[i],a[r[i]]);
        for(int i=2;i<=n;i<<=1)
            for(int j=0;j<n;j+=i)
                for(int k=0;k<i/2;k++)
                {
                    int u=a[j+k],t=(ll)w[op==1?n/i*k:(n-n/i*k)&(n-1)]*a[j+k+i/2]%P;
                    a[j+k]=(u+t)%P,a[j+k+i/2]=(u-t)%P;
                }
        if(op==-1)
        {
            int I=inv(n,P);
            FOR(i,0,n-1)a[i]=(ll)a[i]*I%P;
        }
    }
    void multiply(const int *a,const int *b,int *c,int n1,int n2)
    {
        int n=1;
        while(n<n1+n2-1)n<<=1;
        FOR(i,0,n1-1)A[i]=a[i];
        FOR(i,0,n2-1)B[i]=b[i];
        FOR(i,n1,n-1)A[i]=0;
        FOR(i,n2,n-1)B[i]=0;
        FOR(i,0,n-1)r[i]=(r[i>>1]>>1)|((i&1)*(n>>1));
        w[0]=1,w[1]=Pow(3,(P-1)/n,P);
        FOR(i,2,n-1)w[i]=(ll)w[i-1]*w[1]%P;
        
        DFT(A,1,n),DFT(B,1,n);
        FOR(i,0,n-1)A[i]=(ll)A[i]*B[i]%P;
        DFT(A,-1,n);
        FOR(i,0,n1+n2-2)c[i]=(A[i]+P)%P;
    }
};
int fac[N],ifac[N],f[M];
int A[M],B[M],C[M<<1];
int W[M];
int n,m,s,b;
ll ans;

ll Com(int n,int m){return (ll)fac[n]*ifac[m]%P*ifac[n-m]%P;}

int main()
{
    fac[0]=fac[1]=1;FOR(i,2,N-1)fac[i]=(ll)fac[i-1]*i%P;
    ifac[0]=ifac[1]=1;FOR(i,2,N-1)ifac[i]=(ll)(P-P/i)*ifac[P%i]%P;
    FOR(i,2,N-1)ifac[i]=(ll)ifac[i-1]*ifac[i]%P;
    scanf("%d%d%d",&n,&m,&s);
    FOR(i,0,m)scanf("%d",&W[i]);
    b=min(m,n/s);
    FOR(i,0,b)f[i]=Com(m,i)*Com(n,s*i)%P*fac[s*i]%P*Pow(ifac[s],i,P)%P*Pow(m-i,n-s*i,P)%P;
    FOR(i,0,b)A[i]=(ll)fac[i]*f[i]%P;
    FOR(i,-b,0)B[i+b]=Pow(-1,-i,P)*ifac[-i]%P;
    _Polynomial::multiply(A,B,C,b+1,b+1);
    FOR(i,0,b)(ans+=(ll)W[i]*ifac[i]%P*C[i+b]%P)%=P;
    printf("%lld
",(ans%P+P)%P);
    return 0;
}