[bzoj 3625][Codeforces Round #250]小朋友和二叉树 NTT多项式求逆+多项式开根

Description

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

Input

第一行有2个整数 n,m(1<=n<=10^5; 1<=m<=10^5)。
第二行有n个用空格隔开的互异的整数 c[1],c[2],…,c[n]（1<=c[i]<=10^5)。

Output

输出m行，每行有一个整数。第i行应当含有权值恰为i的神犇二叉树的总数。请输出答案关于998244353(=7*17*2^23+1,一个质数)取模后的结果。

题解

据说是模版题。
考虑最裸的$dp$，设$f[i]$为权值和为$i$时的方案数，然后转移为

$$f[i+j]=\sum c[i] \sum \sum f[j-k]*f[j]$$
仔细盯着式子看可以发现是个卷积套卷积的形式，设$f(x)$为$f$的生成函数，$g(x)$为$c$的生成函数，可以得到 $$f(x)=f(x)^{2}*g(x)+1$$
最后加的系数$1$是为了补齐使它是个生成函数。
有求根公式得
$$f(x)=\frac{1+-\sqrt{1-4g(x)}}{2*g(x)}$$
因为$f(0)=1$
所以

$$f(x)=\frac{1+\sqrt{1-4g(x)}}{2*g(x)}$$
所以 $$f(x)=\frac{2}{1+\sqrt{1-4g(x)}}$$
就可以用NTT，对它求解了。
然后，写了个在 $bzoj$ $T$到死的代码，明明在$cf$上怎么交都能过。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll MOD=998244353;
const int MAXN=2000000;
ll Pinv[MAXN],P[MAXN],TMP[MAXN],QP[MAXN],PS[MAXN];
int R[MAXN],data[MAXN],n,QQQ;
ll read(){ll rt=0;char ch=getchar();while(ch<'0'||ch>'9')ch=getchar();while(ch>='0'&&ch<='9'){rt=rt*10+ch-'0';ch=getchar();}return rt;}
void write(ll num){if(!num)putchar('0');static char str[50];static int cnt=0;while(num){str[++cnt]=num%10+'0';num/=10;}
    while(cnt){putchar(str[cnt--]);}putchar('\n');}
ll mul(ll x,ll y){ll ret=1;while(y){if(y&1){ret=(ret*x)%MOD;}x=(x*x)%MOD;y>>=1;}return ret;}
void NTT(ll *x,int len,int flag){
    for(int i=0;i<len;i++)if(i<R[i])swap(x[i],x[R[i]]);
    for(int i=1;i<len;i<<=1){ll wn=mul(3,(MOD-1)/(i<<1));
        for(int j=0;j<len;j+=(i<<1)){ll w=1;
            for(int k=0;k<i;k++,w=(w*wn)%MOD){
                ll a=x[j+k],b=(x[j+k+i]*w)%MOD;
                x[j+k]=(a+b)%MOD;x[j+k+i]=(a-b+MOD)%MOD;
            }
        }
    }
    if(flag==-1){reverse(x+1,x+len);ll inv=mul(len,MOD-2);for(int i=0;i<len;i++)x[i]=(x[i]*inv)%MOD;}
}
void Out(ll *x,int len){
    for(int i=0;i<len;i++)cout<<x[i]<<" ";
    cout<<endl;
}
void getPinv(int len,ll *x,ll *b){
    if(len==1){b[0]=mul(P[0],MOD-2);return;}
    getPinv((len+1)>>1,x,b);
    int tmp=1,tmpp=len<<1,L=0;
    for(tmp=1;tmp<=tmpp;tmp<<=1)L++;
    for(int i=0;i<tmp;i++)R[i]=0;
    for(int i=0;i<tmp;i++)R[i]=(R[i>>1]>>1)|((i&1)<<(L-1));
    for(int i=0;i<len;i++)x[i]=P[i];
    for(int i=len;i<tmp;i++)x[i]=b[i]=0;
    NTT(x,tmp,233);NTT(b,tmp,233);
    for(int i=0;i<tmp;i++){
        b[i]=b[i]%MOD*(2-(x[i]*b[i])%MOD)%MOD;
        if(b[i]<0)b[i]+=MOD;
    }
    NTT(b,tmp,-1);
    for(int i=len;i<tmp;i++)b[i]=0;
}
void getPsqrt(int len,ll *x,ll *b){
    if(len==1){b[0]=1;return;}
    getPsqrt((len+1)>>1,x,b);
    int tmp=1,tmpp=len<<1,L=0;
    for(tmp=1;tmp<=tmpp;tmp<<=1)L++;
    for(int i=0;i<tmp;i++)Pinv[i]=P[i]=0;
    for(int i=0;i<len;i++)P[i]=(2*b[i])%MOD;
    getPinv(len,TMP,Pinv);
    for(int i=0;i<len;i++){
        x[i]=QP[i];
    }
    for(int i=len;i<tmp;i++)x[i]=0;
    for(int i=0;i<tmp;i++)R[i]=0;
    for(int i=0;i<tmp;i++)R[i]=(R[i>>1]>>1)|((i&1)<<(L-1));
    NTT(x,tmp,233); NTT(Pinv,tmp,233);
    ll Inv = mul(2,MOD-2);
    for(int i=0;i<tmp;i++){
        Pinv[i]=(x[i]*Pinv[i])%MOD;
        if(Pinv[i]<0)Pinv[i]+=MOD;
    }
    NTT(Pinv,tmp,-1);
    for(int i=0;i<tmp;i++)b[i]=(Pinv[i]+(Inv*b[i])%MOD)%MOD;
    for(int i=len;i<tmp;i++)b[i]=0;
}
int main(){
    n=read();int m=0;QQQ=read();
    for(int i=0;i<n;i++){
        int q=read();m=max(q,m);
        data[q]++;
    }
    n=QQQ+1;
    for(int i=0;i<n;i++){
        QP[i]=(MOD-(4*data[i])%MOD)%MOD;
        if(QP[i]<0)QP[i]+=MOD;
    }
    QP[0]++;
    getPsqrt(n,TMP,PS);
    for(int i=0;i<n;i++){
        P[i]=(PS[i])%MOD;
    }
    P[0]++;
    getPinv(n,TMP,Pinv);
    for(int i=1;i<=QQQ;i++){
        write((2*Pinv[i])%MOD);
    }
    return 0;
}