Codeforces 438E The Child and Binary Tree 生成函数+牛顿迭代

题意

考虑一个含有n个互异正整数的序列c[1],c[2],…,c[n]。如果一棵带点权的有根二叉树满足其所有顶点的权值都在集合{c[1],c[2],…,c[n]}中，我们的小朋友就会将其称作神犇的。并且他认为，一棵带点权的树的权值，是其所有顶点权值的总和。
给出一个整数m，你能对于任意的s(1<=s<=m)计算出权值为s的神犇二叉树的个数吗？请参照样例以更好的理解什么样的两棵二叉树会被视为不同的。
我们只需要知道答案关于998244353取模后的值。
$n,m\le 10^5$

分析

这么经典的题居然等到退役之后才来做，真的惭愧。
设$f_i$表示权值为$i$的二叉树数量，那么有$$f_i=\sum f_j{f}_k{g}_{i-j-k}$$
写成生成函数的形式就是$$f(x)=f^2(x)g(x)+1$$
解一下方程可得$$f(x)=\frac{2}{1+\sqrt{1-4g(x)}}$$
直接多项式开根+多项式求逆就好了。

多项式开根和多项式求逆都可以用牛顿迭代来求，具体来说就是当我们要求一个多项式$F(x)$满足
$$G(F(x))\equiv 0(\mathrm{m}\mathrm{o}\mathrm{d}{x}^n)$$
假设我们已经求出了$H(x)$满足
$$G(H(x))\equiv 0(\mathrm{m}\mathrm{o}\mathrm{d}{x}^{\frac{n}{2}})$$
把$G(x)$在$H(x)$处泰勒展开可得
$$G(F(x))=G(H(x))+\frac{G^{\prime }(H(x))}{1!}(F(x)-H(x))+\frac{G^{\prime \prime }(x)}{2!}(F(x)-H(x){)}^2...$$
因为$F(x)$和$H(x)$的前$\frac{n}{2}$项相同，所以
$$G(F(x))\equiv G(H(x))+G^{\prime }(H(x))(F(x)-H(x))(\mathrm{m}\mathrm{o}\mathrm{d}{x}^n)$$
又因为
$$G(F(x))\equiv 0(\mathrm{m}\mathrm{o}\mathrm{d}{x}^n)$$
所以
$$F(x)\equiv H(x)-\frac{G(H(x))}{G^{\prime }(H(x))}(\mathrm{m}\mathrm{o}\mathrm{d}{x}^n)$$
就做完了。

代码

#include
#include
#include
#include
#include

typedef long long LL;

const int N=400005;
const int MOD=998244353;

int n,m,a[N],F[N],G[N],rev[N],b[N],T,lg;

int ksm(int x,int y)
{
	int ans=1;
	while (y)
	{
		if (y&1) ans=(LL)ans*x%MOD;
		x=(LL)x*x%MOD;y>>=1;
	}
	return ans;
}

void NTT(int *a,int f)
{
	for (int i=0;i<T;i++) if (i<rev[i]) std::swap(a[i],a[rev[i]]);
	for (int i=1;i<T;i<<=1)
	{
		int wn=ksm(3,f==1?(MOD-1)/i/2:MOD-1-(MOD-1)/i/2);
		for (int j=0;j<T;j+=(i<<1))
		{
			int w=1;
			for (int k=0;k<i;k++)
			{
				int u=a[j+k],v=(LL)w*a[j+k+i]%MOD;
				a[j+k]=(u+v)%MOD;a[j+k+i]=(u+MOD-v)%MOD;
				w=(LL)w*wn%MOD;
			}
		}
	}
	if (f==-1)
	{
		int ny=ksm(T,MOD-2);
		for (int i=0;i<T;i++) a[i]=(LL)a[i]*ny%MOD;
	}
}

void get_inv(int *a,int n)
{
	if (n==1) {G[0]=ksm(a[0],MOD-2);return;}
	get_inv(a,n/2);
	for (T=1,lg=0;T<=n;T<<=1,lg++);
	for (int i=0;i<T;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(lg-1)),b[i]=0;
	for (int i=0;i<n;i++) b[i]=a[i];
	NTT(b,1);NTT(G,1);
	for (int i=0;i<T;i++) G[i]=(G[i]*2%MOD+MOD-(LL)G[i]*G[i]%MOD*b[i]%MOD)%MOD;
	NTT(G,-1);
	for (int i=n;i<T;i++) G[i]=0;
}

void get_sqrt(int *a,int n)
{
	if (n==1) {F[0]=1;return;}
	get_sqrt(a,n/2);
	get_inv(F,n);
	for (int i=0;i<n;i++) G[i]=(LL)G[i]*((MOD+1)/2)%MOD;
	for (T=1,lg=0;T<n;T<<=1,lg++);
	for (int i=0;i<T;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(lg-1));
	NTT(F,1);
	for (int i=0;i<T;i++) F[i]=(LL)F[i]*F[i]%MOD;
	NTT(F,-1);
	for (int i=0;i<n;i++) (F[i]+=a[i])%=MOD;
	T<<=1;lg++;
	for (int i=0;i<T;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(lg-1));
	NTT(F,1);NTT(G,1);
	for (int i=0;i<T;i++) F[i]=(LL)F[i]*G[i]%MOD,G[i]=0;
	NTT(F,-1);
	for (int i=n;i<T;i++) F[i]=0;
}

int main()
{
	scanf("%d%d",&n,&m);
	for (int i=1;i<=n;i++)
	{
		int x;scanf("%d",&x);
		if (x<=m) a[x]++;
	}
	int k;
	for (k=1;k<=m;k<<=1);
	a[0]=1;
	for (int i=1;i<k;i++) a[i]=(MOD-(LL)a[i]*4%MOD)%MOD;
	get_sqrt(a,k);
	F[0]++;
	get_inv(F,k);
	for (int i=1;i<=m;i++) printf("%d\n",(LL)G[i]*2%MOD);
	return 0;
}