BZOJ 3625 小朋友和二叉树（生成函数+FFT）

Description

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

Input

第一行有$2$个整数$n,m(1\le n\le10^5; 1\le m\le 10^5)$。
第二行有$n$个用空格隔开的互异的整数$c[1],c[2],...,c[n](1\le c[i]\le 10^5)$。

Output

输出$m$行，每行有一个整数。第$i$行应当含有权值恰为$i$的神犇二叉树的总数。请输出答案关于$998244353$取模后的结果。

Sample Input

样例一：

2 3

1 2

样例二：

3 10

9 4 3

样例三：

5 10

13 10 6 4 15

Sample Output

样例一：

1

3

9

样例二：

0

0

1

1

0

2

4

2

6

15

样例三：

0

0

0

1

0

1

0

2

0

5

Solution

设每个点权值的生成函数$A(x)$，设树的生成函数为$F(x)$

考虑有儿子节点和无儿子节点有方程$F(x)=A(x)F^2(x)+1$，解此二次方程得$F(x)=\frac{1\pm \sqrt{1-4A(x)}}{2A(x)}$

注意到$1-4A(x)$开根常数项为$1$，而$A(x)$中没有常数项，故只有将常数项消掉才可以做除法

即解为$F(x)=\frac{1- \sqrt{1-4A(x)}}{2A(x)}$，做多项式开根和逆元即可

Code

#include<cstdio>
#include<iostream>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<vector>
#include<queue>
#include<map>
#include<set>
#include<ctime>
using namespace std;
typedef long long ll;
#define maxn 100005
#define maxfft 262144+5
#define mod 998244353
const double pi=acos(-1.0);
struct cp 
{
    double a,b;
    cp operator +(const cp &o)const {return (cp){a+o.a,b+o.b};}
    cp operator -(const cp &o)const {return (cp){a-o.a,b-o.b};}
    cp operator *(const cp &o)const {return (cp){a*o.a-b*o.b,b*o.a+a*o.b};}
    cp operator *(const double &o)const {return (cp){a*o,b*o};}
    cp operator !() const{return (cp){a,-b};}
}w[maxfft];
int pos[maxfft];
void fft_init(int len)
{
    int j=0;
    while((1<<j)<len)j++;
    j--;
    for(int i=0;i<len;i++)
        pos[i]=pos[i>>1]>>1|((i&1)<<j);
}
void fft(cp *x,int len,int sta)
{
    for(int i=0;i<len;i++)
        if(i<pos[i])swap(x[i],x[pos[i]]);
    w[0]=(cp){1,0};
    for(unsigned i=2;i<=len;i<<=1)
    {
        cp g=(cp){cos(2*pi/i),sin(2*pi/i)*sta};
        for(int j=i>>1;j>=0;j-=2)w[j]=w[j>>1];
        for(int j=1;j<i>>1;j+=2)w[j]=w[j-1]*g;
        for(int j=0;j<len;j+=i)
        {
            cp *a=x+j,*b=a+(i>>1);
            for(int l=0;l<i>>1;l++)
            {
                cp o=b[l]*w[l];
                b[l]=a[l]-o;
                a[l]=a[l]+o;
            }
        }
    }
    if(sta==-1)for(int i=0;i<len;i++)x[i].a/=len,x[i].b/=len;
}
cp x[maxfft],y[maxfft],z[maxfft];
int temp[maxfft];
void FFT(int *a,int *b,int n,int m,int *c)
{
    if(n<=100&&m<=100||min(n,m)<=5)
    {
        for(int i=0;i<n+m-1;i++)temp[i]=0;
        for(int i=0;i<n;i++)
            for(int j=0;j<m;j++)
            {
                temp[i+j]+=(ll)a[i]*b[j]%mod;
                if(temp[i+j]>=mod)temp[i+j]-=mod;
            }
        for(int i=0;i<n+m-1;i++)c[i]=temp[i];
        return ;
    }
    int len=1;
    while(len<n+m)len<<=1;
    fft_init(len);
    for(int i=0;i<len;i++)
    {
        int aa=i<n?a[i]:0,bb=i<m?b[i]:0;
        x[i]=(cp){(aa>>15),(aa&32767)},y[i]=(cp){(bb>>15),(bb&32767)};
    }
    fft(x,len,1),fft(y,len,1);
    for(int i=0;i<len;i++)
    {
        int j=len-1&len-i;
        z[i]=((x[i]+!x[j])*(y[i]-!y[j])+(x[i]-!x[j])*(y[i]+!y[j]))*(cp){0,-0.25};
    }
    fft(z,len,-1);
    for(int i=0;i<n+m-1;i++)
    {
        ll ta=(ll)(z[i].a+0.5)%mod;
        ta=(ta<<15)%mod;
        c[i]=ta;
    }
    for(int i=0;i<len;i++)
    {
        int j=len-1&len-i;
        z[i]=(x[i]-!x[j])*(y[i]-!y[j])*(cp){-0.25,0}+(x[i]+!x[j])*(y[i]+!y[j])*(cp){0,0.25};
    }
    fft(z,len,-1);
    for(int i=0;i<n+m-1;i++)
    {
        ll ta=(ll)(z[i].a+0.5)%mod,tb=(ll)(z[i].b+0.5)%mod;
        ta=(ta+(tb<<30))%mod;
        c[i]=(c[i]+ta)%mod;
    }
}
int inv[maxn];
void init(int n=100001)
{
    inv[1]=1;
    for(int i=2;i<=n;i++)inv[i]=mod-(ll)(mod/i)*inv[mod%i]%mod;
}
int temp1[maxfft],temp2[maxfft],temp3[maxfft],temp4[maxfft];
void Poly_Inv(int *poly,int n,int *ans)
{
    ans[0]=inv[poly[0]];
    for(int i=2;i<=n;i<<=1)
    {
        FFT(poly,ans,i,i/2,temp1);
        FFT(ans,temp1+i/2,i/2,i/2,temp1);
        for(int j=0;j<i/2;j++)ans[j+i/2]=temp1[j]==0?0:mod-temp1[j];
    }
}
void Poly_Log(int *poly,int n,int *ans)
{
    Poly_Inv(poly,n,temp2);
    for(int i=0;i<n-1;i++)ans[i]=(ll)poly[i+1]*(i+1)%mod;
    FFT(ans,temp2,n-1,n,ans);
    for(int i=n-1;i>0;i--)ans[i]=(ll)ans[i-1]*inv[i]%mod;
    ans[0]=0;
}
void Poly_Exp(int *poly,int n,int *ans)
{
    if(n==1)
    {
        ans[0]=1;
        return ;
    }
    Poly_Exp(poly,n/2,ans);
    Poly_Log(ans,n,temp3);
    for(int i=0;i<n;i++)
    {
        temp3[i]=poly[i]-temp3[i];
        if(temp3[i]<0)temp3[i]+=mod;    
    }
    temp3[0]++;
    if(temp3[0]==mod)temp3[0]=0;
    FFT(ans,temp3,n,n,ans);
    for(int i=n;i<2*n;i++)ans[i]=0;
}
void Poly_Root(int *poly,int n,int *ans)
{
    ans[0]=1;
    for(int i=2;i<=n;i<<=1)
    {
        Poly_Inv(ans,i,temp4);
        FFT(ans,ans,i/2,i/2,ans);
        for(int j=0;j<i;j++)ans[j]=(ll)(ans[j]+poly[j])*inv[2]%mod;
        FFT(ans,temp4,i,i,ans);
        for(int j=i;j<2*i;j++)ans[j]=0;
    }
}
int n,m,A[maxfft],B[maxfft];
int main()
{
    init();
    while(~scanf("%d%d",&n,&m))
    {
        memset(A,0,sizeof(A));
        for(int i=1;i<=n;i++)
        {
            int temp;
            scanf("%d",&temp);
            A[temp]=mod-4;
        }
        int len=1;
        while(len<m+1)len<<=1;
        A[0]=1;
        Poly_Root(A,len,B);
        for(int i=1;i<=m;i++)B[i]=(ll)B[i]*inv[2]%mod;
        Poly_Inv(B,len,A);
        for(int i=1;i<=m;i++)printf("%d\n",A[i]);
    }
    return 0;
}