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

 int main(){
     cin>>m>>P>>A>>O>>S>>U;
     for(L i=;i<=m;i++) g[i]=(O*i*i+S*i+U)%P;
     int len=; while(len<=(m*)) len<<=;
     gsum[]=;
     A=min(A,m);
     while(A){

         if(A&){
             NTT(ans,len,); NTT(g,len,);
             for(int i=;i<len;i++) ans[i]=ans[i]*g[i]%MOD;
             NTT(ans,len,-); NTT(g,len,-);
             for(int i=;i<=m;i++)
             ans[i]=(ans[i]+g[i]+gsum[i])%P;
             for(int i=m+;i<len;i++) ans[i]=;
         }
         A>>=;

         g[]++;
         NTT(g,len,); NTT(gsum,len,);
         for(int i=;i<len;i++) gsum[i]=gsum[i]*g[i]%MOD;
         NTT(g,len,-); NTT(gsum,len,-);
         g[]--;
         for(int i=;i<len;i++) if(i>m) gsum[i]=; else gsum[i]%=P;

         NTT(g,len,);
         for(int i=;i<len;i++) g[i]=g[i]*g[i]%MOD;
         NTT(g,len,-);
         for(int i=;i<len;i++) if(i>m) g[i]=; else g[i]%=P;
     }
     cout<<ans[m]<<endl;
 }