【BZOJ3625】【CF438E】小朋友和二叉树 NTT 生成函数 多项式开根 多项式求逆

题目大意

　　考虑一个含有\(n\)个互异正整数的序列\(c_1,c_2,\ldots ,c_n\)。如果一棵带点权的有根二叉树满足其所有顶点的权值都在集合\(\{c_1,c_2,\ldots ,c_n\}\)中，我们的小朋友就会将其称作神犇的。并且他认为，一棵带点权的树的权值，是其所有顶点权值的总和。

　　给出一个整数\(m\)，你能对于任意的\(s(1\leq s\leq m)\)计算出权值为\(s\)的神犇二叉树的个数吗？

　　我们只需要知道答案关于\(998244353\)取模后的值。

　　\(n,m\leq 100000\)

题解

　　设\(a_i\)为\(c\)中有没有\(i\)这个权值。

　　设\(f_i\)为权值和为\(i\)的二叉树个数。

\[\begin{align}
f_0&=1\\
f_i&=\sum_{0\leq j\leq i}a_j\sum_{k=0}^{i-j}f_kf_{i-j-k}\\
\end{align}
\]

　　设

\[\begin{align}
A(x)&=\sum_{i\geq 0}a_ix^i\\
F(x)&=\sum_{i\geq 0}f_ix^i
\end{align}
\]

　　那么

\[\begin{align}
F(x)&=F(x)^2A(x)+1\\
A(x)F(x)^2-F(x)+1&=0\\
\Delta&=1-4A(x)\\
F(x)&=\frac{1\pm\sqrt{1-4A(x)}}{2A(x)}\\
&=\frac{2}{1\pm\sqrt{1-4A(x)}}
\end{align}
\]

　　显然，当常数项有逆元的时候一个多项式才有逆元。而\(a_0=0\)，所以只能取加号。

\[F(x)=\frac{2}{1+\sqrt{1-4A(x)}}
\]

　　直接多项式开根+多项式求逆。

　　时间复杂度：\(O(n\log n)\)

代码

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cstdlib>
#include<ctime>
#include<utility>
#include<cmath>
#include<functional>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
void sort(int &a,int &b)
{
	if(a>b)
		swap(a,b);
}
void open(const char *s)
{
#ifndef ONLINE_JUDGE
	char str[100];
	sprintf(str,"%s.in",s);
	freopen(str,"r",stdin);
	sprintf(str,"%s.out",s);
	freopen(str,"w",stdout);
#endif
}
int rd()
{
	int s=0,c;
	while((c=getchar())<'0'||c>'9');
	do
	{
		s=s*10+c-'0';
	}
	while((c=getchar())>='0'&&c<='9');
	return s;
}
int upmin(int &a,int b)
{
	if(b<a)
	{
		a=b;
		return 1;
	}
	return 0;
}
int upmax(int &a,int b)
{
	if(b>a)
	{
		a=b;
		return 1;
	}
	return 0;
}
const ll p=998244353;
const ll g=3;
ll fp(ll a,ll b)
{
    ll s=1;
    while(b)
    {
        if(b&1)
            s=s*a%p;
        a=a*a%p;
        b>>=1;
    }
    return s;
}
const int maxn=600000;
ll inv[maxn];
namespace ntt
{
    ll w1[maxn];
    ll w2[maxn];
    int rev[maxn];
    int n;
    void init(int m)
    {
        n=1;
        while(n<m)
            n<<=1;
        int i;
        for(i=2;i<=n;i<<=1)
        {
            w1[i]=fp(g,(p-1)/i);
            w2[i]=fp(w1[i],p-2);
        }
        rev[0]=0;
        for(i=1;i<n;i++)
            rev[i]=(rev[i>>1]>>1)|((i&1)*(n>>1));
    }
    void ntt(int *a,int t)
    {
        int i,j,k;
        int u,v,w,wn;
        for(i=0;i<n;i++)
            if(rev[i]<i)
                swap(a[i],a[rev[i]]);
        for(i=2;i<=n;i<<=1)
        {
            wn=(t==1?w1[i]:w2[i]);
            for(j=0;j<n;j+=i)
            {
                w=1;
                for(k=j;k<j+i/2;k++)
                {
                    u=a[k];
                    v=1ll*a[k+i/2]*w%p;
					a[k]=(u+v)%p;
					a[k+i/2]=(u-v)%p;
                    w=1ll*w*wn%p;
                }
            }
        }
        if(t==-1)
        {
            u=fp(n,p-2);
            for(i=0;i<n;i++)
                a[i]=1ll*a[i]*u%p;
        }
    }
    int x[maxn];
    int y[maxn];
    int z[maxn];
    void copy_clear(int *a,int *b,int m)
    {
        int i;
        for(i=0;i<m;i++)
            a[i]=b[i];
        for(i=m;i<n;i++)
            a[i]=0;
    }
    void copy(int *a,int *b,int m)
    {
        int i;
        for(i=0;i<m;i++)
            a[i]=b[i];
    }
    void inverse(int *a,int *b,int m)
    {
        if(m==1)
        {
            b[0]=fp(a[0],p-2);
            return;
        }
        inverse(a,b,m>>1);
        init(m<<1);
        copy_clear(x,a,m);
        copy_clear(y,b,m>>1);
        ntt(x,1);
        ntt(y,1);
        int i;
        for(i=0;i<n;i++)
            x[i]=y[i]*(2-1ll*x[i]*y[i]%p)%p;
    	ntt(x,-1);
    	copy(b,x,m);
    }
    int c[maxn],d[maxn];
    void sqrt(int *a,int *b,int m)
    {
    	if(m==1)
    	{
    		if(a[0]==1)
    			b[0]=1;
    		else if(a[0]==0)
    			b[0]=0;
    		else
    			//我也不会
				;
			return;
		}
		sqrt(a,b,m>>1);
//		copy_clear(c,b,m>>1);
		int i;
		for(i=m;i<m<<1;i++)
			b[i]=0;
		inverse(b,d,m);
		init(m<<1);
		for(i=m;i<m<<1;i++)
			b[i]=d[i]=0;
		ll inv2=fp(2,p-2);
		copy_clear(x,a,m);
		ntt(x,1);
		ntt(d,1);
		for(i=0;i<n;i++)
			x[i]=1ll*x[i]*d[i]%p;
		ntt(x,-1);
		for(i=0;i<m;i++)
			b[i]=(1ll*(b[i]+x[i])%p*inv2)%p;
	}
};
int a[maxn];
int b[maxn];
int main()
{
	open("bzoj3625");
	int n,m;
	int i;
	int x;
	scanf("%d%d",&n,&m);
	for(i=1;i<=n;i++)
	{
		scanf("%d",&x);
		a[x]=1;
	}
	int k=1;
	while(k<=m)
		k<<=1;
	for(i=0;i<k;i++)
		a[i]=-4ll*a[i]%p;
	a[0]=(a[0]+1)%p;
	ntt::sqrt(a,b,k);
	b[0]=(b[0]+1)%p;
	ntt::inverse(b,a,k);
	for(i=1;i<=m;i++)
	{
		a[i]=(a[i]*2ll%p+p)%p;
		printf("%d\n",a[i]);
	}
	return 0;
}