【bzoj4332】【JSOI2012】 分零食 生成函数 FFT

我们构造$f(x)$的生成函数$G(x)$，那么显然$[x^k]G(x)=Ok^2+Sk+U$

那么显然，答案即为$sum_{i=1}^{n} [x^m]G^i(x)$

我们构造答案的生成函数$F(x)=sum_{i=1}^{n} G^i(x)$

根据等比数列求和公式，$F(x)=G(x)dfrac{1-G^{A}(x)}{1-G(x)}$

如果去等比数列求和的话，你需要多项式快速幂+多项式求逆，时间复杂度显然是$O(m log m)$的。

然而这个模数并不是质数，所以这么搞不是很好搞。

我们可以用一个类似快速幂的方式，去算出$sum_{i-1}^{2^k-1}G^i(x)$的值。

这么搞的时间复杂度显然是$O(m log m log A)$。

然后就没了

第一次自己推出生成函数的题美滋滋

#include<bits/stdc++.h>
#define MOD 998244353
#define L long long
#define M 1<<15
#define G 3
using namespace std;

L pow_mod(L x,L k){
    L ans=1;
    while(k){
        if(k&1) ans=ans*x%MOD;
        x=x*x%MOD; k>>=1;
    }
    return ans;
}
void change(L a[],int n){
    for(int i=0,j=0;i<n-1;i++){
        if(i<j) swap(a[i],a[j]);
        int k=n>>1;
        while(j>=k) j-=k,k>>=1;
        j+=k;
    }
}
void NTT(L a[],int n,int on){
    change(a,n);
    for(int h=2;h<=n;h<<=1){
        L wn=pow_mod(G,(MOD-1)/h);
        for(int j=0;j<n;j+=h){
            L w=1;
            for(int k=j;k<j+(h>>1);k++){
                L u=a[k],t=w*a[k+(h>>1)]%MOD;
                a[k]=(u+t)%MOD; 
                a[k+(h>>1)]=(u-t+MOD)%MOD;
                w=w*wn%MOD;
            }
        }
    }
    if(on==-1){
        L inv=pow_mod(n,MOD-2);
        for(int i=0;i<n;i++) a[i]=a[i]*inv%MOD;
        reverse(a+1,a+n);
    }
}
L m,P,A,O,S,U;
L g[M]={0},gsum[M]={0},ans[M]={0};

int main(){
    cin>>m>>P>>A>>O>>S>>U;
    for(L i=1;i<=m;i++) g[i]=(O*i*i+S*i+U)%P;
    int len=1; while(len<=(m*2)) len<<=1;
    gsum[0]=1;
    A=min(A,m);
    while(A){
        
        if(A&1){
            NTT(ans,len,1); NTT(g,len,1);
            for(int i=0;i<len;i++) ans[i]=ans[i]*g[i]%MOD;
            NTT(ans,len,-1); NTT(g,len,-1);
            for(int i=1;i<=m;i++) 
            ans[i]=(ans[i]+g[i]+gsum[i])%P;
            for(int i=m+1;i<len;i++) ans[i]=0; 
        }
        A>>=1;
        
        g[0]++;
        NTT(g,len,1); NTT(gsum,len,1);
        for(int i=0;i<len;i++) gsum[i]=gsum[i]*g[i]%MOD;
        NTT(g,len,-1); NTT(gsum,len,-1);
        g[0]--;
        for(int i=0;i<len;i++) if(i>m) gsum[i]=0; else gsum[i]%=P;
        
        NTT(g,len,1); 
        for(int i=0;i<len;i++) g[i]=g[i]*g[i]%MOD;
        NTT(g,len,-1);
        for(int i=0;i<len;i++) if(i>m) g[i]=0; else g[i]%=P;
    }
    cout<<ans[m]<<endl;
}