[LOJ3106][TJOI2019]唱、跳、rap和篮球:DP+生成函数+NTT+容斥原理

分析

\(f(i)\)表示共\(i\)组同学讨论cxk的位置的方案数(不考虑其他位置上的人的爱好),这个数组可以很容易地通过依次考虑每个位置是否是四个人中最后一个人的位置来递推求解,时间复杂度\(O(n^2)\)

\(g(i)\)表示共\(i\)组同学讨论cxk,剩下的\(n-4i\)个位置上的人的爱好的方案数。这个数组可以通过对每种情况,分别写出四种爱好的\(EGF\),然后\(NTT\)合并求解,时间复杂度\(O(n^2 \log n)\)

统计答案的话很简单,容斥一下就好了:

\[ans=\sum_{i=0}^{limit}(-1)^if(i)g(i)\]

代码

#include <bits/stdc++.h>

#define rin(i,a,b) for(int i=(a);i<=(b);++i)
#define irin(i,a,b) for(int i=(a);i>=(b);--i)
#define trav(i,a) for(int i=head[a];i;i=e[i].nxt)
#define Size(a) (int)a.size()
#define pb push_back
#define mkpr std::make_pair
#define fi first
#define se second
#define lowbit(a) ((a)&(-(a)))
typedef long long LL;

using std::cerr;
using std::endl;

inline int read(){
    int x=0,f=1;char ch=getchar();
    while(!isdigit(ch)){if(ch=='-')f=-1;ch=getchar();}
    while(isdigit(ch)){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}

const int MAXN=1005;
const int NTT=1048576;
const int MOD=998244353;
const int G=3,INVG=332748118;

int N,a,b,c,d;
int n,m,len,rev[MAXN<<2];
int w[NTT+5],iw[NTT+5]; 
int fac[MAXN],invf[MAXN];
int f[MAXN][MAXN],g[MAXN];
int A[MAXN<<2],B[MAXN<<2],C[MAXN<<2],D[MAXN<<2];

inline int add(int x,int y){
    return x+y<MOD?x+y:x+y-MOD;
}

inline int mns(int x,int y){
    return x-y>=0?x-y:x-y+MOD;
}

inline int qpow(int x,int y){
    int ret=1,tt=x%MOD;
    while(y){
        if(y&1)ret=1ll*ret*tt%MOD;
        tt=1ll*tt*tt%MOD;
        y>>=1;
    }
    return ret;
}

void ntt(int *c,int dft){
    rin(i,0,n-1)if(i<rev[i])std::swap(c[i],c[rev[i]]);
    for(int mid=1;mid<n;mid<<=1){
        int r=(mid<<1),u=NTT/r;
        for(int l=0;l<n;l+=r){
            for(int i=0,v=0;i<mid;++i,v+=u){
                int x=c[l+i],y=1ll*c[l+mid+i]*(dft>0?w[v]:iw[v])%MOD;
                c[l+i]=add(x,y);
                c[l+mid+i]=mns(x,y);
            }
        }
    }
    if(dft<0){
        int invn=qpow(n,MOD-2);
        rin(i,0,n-1)c[i]=1ll*c[i]*invn%MOD;
    }
}

void init(){
    fac[0]=1;rin(i,1,N)fac[i]=1ll*fac[i-1]*i%MOD;
    invf[N]=qpow(fac[N],MOD-2);irin(i,N-1,0)invf[i]=1ll*invf[i+1]*(i+1)%MOD;
    w[0]=iw[0]=1;w[1]=qpow(G,(MOD-1)/NTT);iw[1]=qpow(INVG,(MOD-1)/NTT);
    rin(i,2,NTT-1)w[i]=1ll*w[i-1]*w[1]%MOD,iw[i]=1ll*iw[i-1]*iw[1]%MOD;
}

void prepare(){
    for(n=1,len=0;n<=m;n<<=1,++len);
    rin(i,1,n-1)rev[i]=((rev[i>>1]>>1)|((i&1)<<(len-1)));
}

int main(){
    N=read(),a=read(),b=read(),c=read(),d=read();
    init();
    int lim=std::min(N/4,std::min(std::min(a,b),std::min(c,d)));
    f[0][0]=1;
    rin(i,1,N){
        rin(j,0,lim){
            f[i][j]=f[i-1][j];
            if(j&&i>3)f[i][j]=add(f[i][j],f[i-4][j-1]);
        }
    }
    rin(i,0,lim){
        m=a-i+b-i+c-i+d-i;prepare();
        rin(j,0,a-i)A[j]=invf[j];
        rin(j,0,b-i)B[j]=invf[j];
        rin(j,0,c-i)C[j]=invf[j];
        rin(j,0,d-i)D[j]=invf[j];
        ntt(A,1);ntt(B,1);ntt(C,1);ntt(D,1);
        rin(j,0,n-1)A[j]=1ll*A[j]*B[j]%MOD*C[j]%MOD*D[j]%MOD;;
        ntt(A,-1);
        g[i]=1ll*A[N-i*4]*fac[N-i*4]%MOD;
        memset(A,0,n*sizeof(int));
        memset(B,0,n*sizeof(int));
        memset(C,0,n*sizeof(int));
        memset(D,0,n*sizeof(int));
    }
    int ans=0,sgn=MOD-1;
    rin(i,0,lim){
        sgn=MOD-sgn;
        ans=(ans+1ll*sgn*f[N][i]%MOD*g[i])%MOD;
    }
    printf("%d\n",ans);
    return 0;
}

posted on 2019-05-05 11:12 ErkkiErkko 阅读(...) 评论(...) 编辑 收藏

统计