JSOI2012 分零食 生成函数 FFT

Posted 2020-12-25 $mathit{AlphaINF}$

我们构造$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)=dfrac{1-G(x)}{1-G^{A+1}(x)}$

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

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

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

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

然后就没了

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

 1 #include<bits/stdc++.h>
 2 #define MOD 998244353
 3 #define L long long
 4 #define M 1<<15
 5 #define G 3
 6 using namespace std;
 7 
 8 L pow_mod(L x,L k){
 9     L ans=1;
10     while(k){
11         if(k&1) ans=ans*x%MOD;
12         x=x*x%MOD; k>>=1;
13     }
14     return ans;
15 }
16 void change(L a[],int n){
17     for(int i=0,j=0;i<n-1;i++){
18         if(i<j) swap(a[i],a[j]);
19         int k=n>>1;
20         while(j>=k) j-=k,k>>=1;
21         j+=k;
22     }
23 }
24 void NTT(L a[],int n,int on){
25     change(a,n);
26     for(int h=2;h<=n;h<<=1){
27         L wn=pow_mod(G,(MOD-1)/h);
28         for(int j=0;j<n;j+=h){
29             L w=1;
30             for(int k=j;k<j+(h>>1);k++){
31                 L u=a[k],t=w*a[k+(h>>1)]%MOD;
32                 a[k]=(u+t)%MOD; 
33                 a[k+(h>>1)]=(u-t+MOD)%MOD;
34                 w=w*wn%MOD;
35             }
36         }
37     }
38     if(on==-1){
39         L inv=pow_mod(n,MOD-2);
40         for(int i=0;i<n;i++) a[i]=a[i]*inv%MOD;
41         reverse(a+1,a+n);
42     }
43 }
44 L m,P,A,O,S,U;
45 L g[M]={0},gsum[M]={0},ans[M]={0};
46 
47 int main(){
48     cin>>m>>P>>A>>O>>S>>U;
49     for(L i=1;i<=m;i++) g[i]=(O*i*i+S*i+U)%P;
50     int len=1; while(len<=(m*2)) len<<=1;
51     gsum[0]=1;
52     A=min(A,m);
53     while(A){
54         
55         if(A&1){
56             NTT(ans,len,1); NTT(g,len,1);
57             for(int i=0;i<len;i++) ans[i]=ans[i]*g[i]%MOD;
58             NTT(ans,len,-1); NTT(g,len,-1);
59             for(int i=1;i<=m;i++) 
60             ans[i]=(ans[i]+g[i]+gsum[i])%P;
61             for(int i=m+1;i<len;i++) ans[i]=0; 
62         }
63         A>>=1;
64         
65         g[0]++;
66         NTT(g,len,1); NTT(gsum,len,1);
67         for(int i=0;i<len;i++) gsum[i]=gsum[i]*g[i]%MOD;
68         NTT(g,len,-1); NTT(gsum,len,-1);
69         g[0]--;
70         for(int i=0;i<len;i++) if(i>m) gsum[i]=0; else gsum[i]%=P;
71         
72         NTT(g,len,1); 
73         for(int i=0;i<len;i++) g[i]=g[i]*g[i]%MOD;
74         NTT(g,len,-1);
75         for(int i=0;i<len;i++) if(i>m) g[i]=0; else g[i]%=P;
76     }
77     cout<<ans[m]<<endl;
78 }