XX Open Cup GP of Gomel B (Baekjoon18743) Bin【生成函数+分治fft】

XX Open Cup GP of Gomel B (Baekjoon18743) Bin

$f(n)$ 表示叶子节点为 $n$ 的方案数。则根据定义：
$$f(n)=\sum _{i\le n-i+k}f(i)\times f(n-i)$$
可以发现因为 $k$ 很小，所以 $i$ 是枚举了 $1$ 到 $n$ 的一半多一点点的位置。

如果让 $i$ 枚举 $1$ 到 $n$ ，新的和就变成了 $f(n)$ 的两倍少中间那一点的样子。

于是考虑两边同时 $\times 2$ 把限制放松
$$2f(n)=\sum _{1\le i<n}f(i)\times f(n-i)+\sum _{|i-(n-i)|\le k}f(i)f(n-i)$$
把它看成多项式，那就是卷积形式。

左边的和式就是分治fft可做，（注意是自己卷自己和分治fft的模板不太一样。$O(n{\log }^{2}n)$ 。

右边的和式在分治到单个节点的时候暴力算，$O(nk)$ 。

时间复杂度 $O(n{\log }^{2}n+nk)$ 。

#include <bits/stdc++.h>
#define N 2000006
using namespace std;
typedef long long ll; 
typedef vector<ll> vec;
const ll mod=998244353;
int now_limit;
int rev[N];
ll wq[20][N],fac[N],inv[N];
ll ksm(ll x,ll y){
	ll res=1; while(y){ if(y&1)res=res*x%mod; x=x*x%mod; y>>=1; }
	return res;
}
void init_NTT(int limit){
	if(limit==now_limit) return;
	now_limit=limit; int l=0; while((1<<l)<limit)l++;
	for(int i=0;i<limit;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(l-1));
	if(wq[0][0]==0){
		while(limit<1000000) limit<<=1;
		for(int mid=1,l=0;mid<limit;mid<<=1,l++){
			wq[l][0]=1,wq[l][1]=ksm(3,(mod-1)/(mid<<1));
			for(int k=2;k<mid;k++) wq[l][k]=wq[l][k-1]*wq[l][1]%mod;
		}
	}
}
void NTT(vec &a,int limit,int flag){
	init_NTT(limit);
	while(a.size()<limit) a.push_back(0);
	for(int i=0;i<limit;i++) if(i<rev[i])swap(a[i],a[rev[i]]);
	ll x,y;
	for(int mid=1,l=0;mid<limit;mid<<=1,l++)
		for(int j=0;j<limit;j+=(mid<<1))
			for(int k=0;k<mid;k++){
				x=a[j+k],y=a[j+k+mid]*wq[l][k]%mod;
				a[j+k]=(x+y)%mod,a[j+k+mid]=(x-y+mod)%mod;
			}
	if(flag==-1){
		x=ksm(limit,mod-2);
		for(int i=0;i<limit;i++) a[i]=a[i]*x%mod;
		reverse(&a[1],&a[limit]);
	}
}
vec g,f;
int n,k,inv2;
void cdq_NTT(int l,int r){
	if(r-l<2){
		if(l>k)
		for(int i=(l-k+1)/2;i+i<l;i++) f[l]=(f[l]+f[i]*f[l-i]%mod)%mod;
		int m=min((l+k)/2,l);
		for(int i=(l+1)/2;i<=m;i++) f[l]=(f[l]+f[i]*f[l-i]%mod)%mod;
		f[l]=f[l]*inv2%mod;
		return;
	}
	int mid=(l+r+1)>>1;
	cdq_NTT(l,mid);
	vec a,b; a=vector<ll>(&f[l],&f[mid]); b=vector<ll>(&f[0],&f[min(r-l,mid)]);
	int limit=1; while(limit<r-l) limit<<=1;
	NTT(a,limit,1),NTT(b,limit,1);
	for(int i=0;i<limit;i++) a[i]=a[i]*b[i]%mod;
	NTT(a,limit,-1);
	if(l!=0) for(int i=mid;i<r;i++) f[i]=(f[i]+a[i-l]+a[i-l])%mod;
	else for(int i=mid;i<r;i++) f[i]=(f[i]+a[i-l])%mod;
	cdq_NTT(mid,r);
}
int main(){
//	freopen("test.in","r",stdin);
	cin>>n>>k;
	f.resize(n+1);
	f[1]=2; inv2=ksm(2,mod-2);
	cdq_NTT(0,n+1);
	cout<<f[n]<<'\n';
}