bzoj3625-小朋友和二叉树

题目

给出(n)个互异的正整数序列，(c_1,c_2...c_n)。定义一颗带点权树的权值为所有点的权值和。给出一个整数(m)，请计算对于所有的(i=1,2,3...,m)，用(c)中的数能组成多少不同的二叉树（每个数可以重复使用）且二叉树的权值为(i)。答案对(7*17*2^{23}+1=998244353)取模。

Sample Input

对于(n=2，c={1,2}，i=3)，有九个这样的二叉树：

分析

其实看到那个模数就可以想到要用到(NTT)了。
可以说第一次遇到这种题，但这种生成函数法的思想以前有接触过。
两个二叉树不同就是不全等，所以我们可以分成三部分，当前点，左子树和右子树。我们设一个函数(f)，(f_i)表示一个子树的点权和为(i)的情况数。再设(v_i)为权值(i)的个数，即有或没有。那么有：

[egin{aligned} f_0&=1 \ f_k&=sum _{i=0}^kv_isum _{j=0}^{k-i}f_jf_{k-i-j} end{aligned} ]

这个式子其实可以用一个生成函数来表示。生成函数是指一种关于另一个函数(g)的形式函数，它的(x^i)项的系数为(g(i))，可以这样表示：

[G(x)=sum _{i=0}^{infty}f(i)x^i ]

我们设(F)为(f)的生成函数，(C)为(v)的生成函数，那么有：

[F(x)=C(x)*F(x)*F(x)+1 ]

那个1是从(f_0=1)来的。
我们要求的是(F)的(1-m)项的系数，所以要把(F)解出来。

[egin{aligned} F(x)=C(x)*F(x)*F(x)+1 \ F(x)=frac{1pmsqrt{1-4C(x)}}{2C(x)} end{aligned} ]

通过F(0)=1判断取值，若取正号：

[egin{aligned} lim _{x o 0}F(x)&=frac{1+sqrt{1-4C(x)}}{2C(x)} \ &=infty end{aligned} ]

矛盾，所以不可取。若取负号：

[egin{aligned} lim _{x o 0}F(x)&=frac{1-sqrt{1-4C(x)}}{2C(x)} \ lim _{x o 0}F(x)&=frac{4C(x)}{2C(x)(1+sqrt {1-4C(x)})} \ lim _{x o 0}F(x)&=frac{2}{1+sqrt {1-4C(x)}} \ &=1 end{aligned} ]

所以我们要计算的(F)就是：

[egin{aligned} F(x)&=frac{1-sqrt{1-4C(x)}}{2C(x)} \ F(x)&=frac{2}{1+sqrt {1-4C(x)}} \ end{aligned} ]

这里涉及到了多项式模意义下的求逆与开根。

多项式开根

运用分治和(NTT)即可在(O(nlog^2n))的时间内实现。
考虑边界情况，如果多项式只有常数项，那么它的开根就是常数的开根。接下来考虑分治。
我们记符号([k])表示多项式的(0...k-1)项模意义下同余。假设我们已经求出了([frac{n}{2}])的答案，现在考虑怎样求([n])的。

[egin{aligned} 对F(x)开根 已知：G(x)^2=F(x) &&[frac{n}{2}] \ 求：H(x)^2=F(x) &&[n] \ H(x)^2-G(x)^2=0 &&[frac{n}{2}] \ (H(x)^2-G(x)^2)^2=0 &&[n] \ (H(x)^2+G(x)^2)^2=4H(x)^2G(x)^2 &&[n] \ H(x)=frac{H(x)^2+G(x)^2}{2G(x)} &&[n] \ H(x)=frac{F(x)+G(x)^2}{2G(x)} &&[n] \ H(x)=frac{F(x)}{2G(x)}+frac{G(x)}{2} &&[n] \ end{aligned} ]

第二行，由于(H(x))满足(0...n-2)都符合要求，那么它一定满足(0..frac{n}{2}-1)项的要求。
这里需要解释一下为什么第四行处平方后模数也跟着平方。由于第三行两个函数平方相减为0，所以对于所有的(x)都成立，所以只能是系数全为0。因此，后面在平方的时候，对于(0...n-1)项，它们得出的时候必定会有之前的(0...{frac{n}{2}-1})的因数，所以平方后(0...n-1)项也全为0。
多项式可以开根的条件就是常数项在模意义下可以开根，所以判断起来会比较麻烦，还要用到离散对数什么的。

多项式求逆

同样思路，我们考虑边界情况，如果多项式只有常数项，那么那么它的逆就是常数项的逆元。同样考虑分治，我们已经求出了能满足前一半的结果，现在考虑怎样求能够满足整个区间的结果。

[egin{aligned} 对F(x)求逆 \ 已知：G(x)F(x)=1 &&[frac{n}{2}] \ 求：H(x)F(x)=1 &&[n] \ (H(x)-G(x))F(x)=0 &&[frac{n}{2}] \ ecause F(x) e 0 \ herefore H(x)-G(x)=0 &&[frac{n}{2}] \ herefore (H(x)-G(x))^2=0 &&[n] \ H(x)^2-2H(x)G(x)+G(x)^2=0 & &[n] \ 移项，两边同时乘以F(x) \ H(x)=2G(x)-F(x)G(x)^2 &&[n] end{aligned} ]

就这样直接递归计算一下求逆和开根即可，也可以写成非递归的。

代码

调了好久，主要在于之前写的时候每次传入的(len)没有想清楚要不要左移右移，导致了一大堆错误。其实这种情况下重新写会更好。
小视野上两秒一个点跑过去了，bzoj不行，codeforces过了。

#include<cstdio>
#include<algorithm>
#include<cctype>
#include<cmath>
#include<cstring>
using namespace std;
typedef long long giant;
giant read() {
	giant x=0,f=1;
	char c=getchar();
	for (;!isdigit(c);c=getchar()) if (c=='-') f=-1;
	for (;isdigit(c);c=getchar()) x=x*10+c-'0';
	return x*f;
}
const giant q=998244353;
const giant g=3;
const giant maxj=19;
const giant maxn=(1<<maxj)+1;
giant c[maxn],xj,go[maxn];
giant ws[maxj+1][2];
giant mi(giant x,giant y) {
	giant ret=1;
	while (y) {
		if (y&1) (ret*=x)%=q;
		y>>=1,(x*=x)%=q;
	}
	return ret;
}
giant inv(giant x) {
	return mi(x,q-2);
}
void ntt(giant a[],giant M,giant o=1) {
	go[0]=0;
	for (xj=0;(1<<xj)!=M;++xj);
	for (giant i=0;i<M;++i) go[i]=(go[i>>1]>>1)|((i&1)<<(xj-1));
	for (giant i=0;i<M;++i) if (i<go[i]) swap(a[i],a[go[i]]);
	for (giant i=2,th=1;i<=M;i<<=1,++th) {
		for (giant j=0;j<M;j+=i) {
			giant w=1;
			for (giant k=j;k<j+i/2;++k,(w*=ws[th][o])%=q) {
				giant u=a[k]%q,v=(a[k+i/2]*w)%q;
				a[k]=(u+v+q)%q,a[k+i/2]=((u-v+q)%q+q)%q;
			}
		} 	
	}
	if (!o) {
		giant tmp=inv(M);
		for (giant i=0;i<M;++i) (a[i]*=tmp)%=q;
	}
}
giant g1[maxn],g2[maxn],f[maxn],ia[maxn],ra[maxn],ig1[maxn],ig2[maxn],iif[maxn];
void Inv(giant a[],giant len) {
	if (len==1) {
		ia[0]=inv(a[0]);
		return; 
	}
	Inv(a,len>>1);
	for (giant i=0;i<(len>>1);++i) ig1[i]=ig2[i]=ia[i];
	for (giant i=0;i<len;++i) iif[i]=a[i];
	ntt(iif,len<<1),ntt(ig2,len<<1),ntt(ig1,len<<1); //here
	for (giant i=0;i<(len<<1);++i) {
		iif[i]=((2*ig1[i]%q-(ig2[i]*ig2[i]%q)*iif[i]%q)%q+q)%q; //here
	}
	ntt(iif,len<<1,0); //here
	for (giant i=0;i<len;++i) ia[i]=iif[i];
	for (giant i=0;i<(len<<1);++i) iif[i]=ig1[i]=ig2[i]=0; //here
}
void Root(giant a[],giant len) { 
	if (len==1) {
		ra[0]=sqrt(a[0]);
		return;
	}
	Root(a,len>>1);
	giant it=inv(2);
	for (giant i=0;i<len;++i) f[i]=a[i];
	for (giant i=0;i<(len>>1);++i) {
		g1[i]=(ra[i]*it)%q;
		g2[i]=(ra[i]*2)%q;
	}
	Inv(g2,len);
	ntt(f,len<<1),ntt(ia,len<<1); // here
	for (giant i=0;i<(len<<1);++i) (f[i]*=ia[i])%=q;
	ntt(f,len<<1,0); // here
	for (giant i=0;i<len;++i) ra[i]=(g1[i]+f[i])%q;
	for (giant i=0;i<(len<<1);++i) f[i]=g1[i]=g2[i]=ia[i]=0; //here
}
int main() {
	#ifndef ONLINE_JUDGE
		freopen("test.in","r",stdin);
	#endif
	giant n=read(),m=read();
	giant tmp=1;
	while (tmp<=(m<<1)) tmp<<=1;
	for (giant i=1;i<=n;++i) c[read()]++;
	for (giant i=0;i<tmp;++i) c[i]=((c[i]*-4)%q+q)%q;
	for (giant i=0;i<=maxj;++i) {
		ws[i][1]=mi(g,(q-1)/(1<<i));
		ws[i][0]=inv(ws[i][1]);
	}
	(c[0]+=1)%=q;
	Root(c,tmp>>1);
	(ra[0]+=1)%=q;
	Inv(ra,tmp);
	for (giant i=1;i<=m;++i) (ia[i]*=2)%=q;
	for (giant i=1;i<=m;++i) printf("%lld
",ia[i]);
}