Codeforces1096G Lucky Tickets(NTT优化dp)

Posted 2021-02-07 psychicboom

设(f[i][j])表示填了(i)个数，数位和为(j)的方案数

于是方程为:
[f[i][j]=sum_{k=1}^9 f[i-1][j-k]*[CanUse[k]==1]]

其中(CanUse[i])表示是否可用(i)这个数字

最终答案为：
[sum_{i=1}^{9*(n/2)}f[n/2][j]]

直接转移肯定(T)飞，需要一些优化。于是我们观察到这个式子是卷积形式的式子，直接上(NTT)板子+快速幂即可

(P.S.)一些优化

1. 可以把每个(pow(g,(mod-1)/i))和(pow(g,-(mod-1)/i))预处理出来，(NTT)过程中直接调用

2. 快速幂前可以直接把系数表示化为点值表示，快速幂时直接乘，快速幂后再换回去系数表示

所以时间复杂度为(0(9*(n/2)log_29*(n/2)))

代码：

#include <bits/stdc++.h>
#define ll long long
#define mod 998244353ll
#define N 1048580
using namespace std;
ll G[N],F[N],tmp[N],tot[N],invG[N],inv,ans=0;
const ll g=3;
int rev[N],len,x;
void NTT(ll *a,int n,int f){
    int i,j,k;
    for(i=0;i<n;++i) if(i<rev[i]) swap(a[i],a[rev[i]]);
    for(i=2;i<=n;i<<=1){
        ll omega;
        int m=i>>1;
        (f==1)?omega=G[i]:omega=invG[i];
        for(j=0;j<n;j+=i){
            ll p=1;
            for(k=j;k<j+m;++k,(p*=omega)%=mod){
                ll u=a[k],v=p*a[k+m]%mod;
                a[k]=(u+v)%mod,a[k+m]=(u-v+mod)%mod;
            }
        }
    } 
    if(f==-1) for(i=0;i<n;++i) (a[i]*=inv)%=mod; 
}       
ll ksm(ll x,int y){
    ll ans=1;
    while(y){
        if(y&1) (ans*=x)%=mod;
        (x*=x)%=mod;
        y>>=1;
    }
    return ans;
}
void init(int n){
    int i;
    for(len=1,x=0;len<=n;len<<=1,++x);
    for(i=0;i<len;++i) rev[i]=(rev[i>>1]>>1)|((i&1)<<(x-1));
    G[len]=ksm(g,(mod-1)/len),invG[len]=ksm(g,mod-1-(mod-1)/len);
    for(i=(len>>1);i>=2;i>>=1){
        G[i]=G[i<<1]*G[i<<1]%mod;
        invG[i]=invG[i<<1]*invG[i<<1]%mod;
    }   
}
void mul(ll *a,ll *b){
    for(int i=0;i<len;++i) a[i]=a[i]*b[i]%mod;
}
int main(){
    int n,d,i;
    scanf("%d%d",&n,&d);n/=2;
    for(i=1;i<=d;++i) scanf("%d",&x),F[x]=1;
    init(9*n);
    tot[0]=1;inv=ksm(1ll*len,mod-2);
    NTT(F,len,1),NTT(tot,len,1);
    x=n;
    while(x){
        if(x&1) mul(tot,F);
        x>>=1;
        if(x) mul(F,F);
    }
    NTT(tot,len,-1);        
    for(i=0;i<=9*n;++i) (ans+=tot[i]*tot[i]%mod)%=mod;
    cout<<ans;
}