【模板】多项式开方

蒟蒻写题解实在不易

前置芝士

NTT与多项式求逆

推式

推式中如有不理解的地方在多项式求逆的题解中均有详细说明
求\(B(x)\)，使得\(B(x)^2\equiv A(x)(mod x^n)\)

\[\begin{aligned}\\ B(x)^2\equiv A(x)(mod x^n),B(x)^2&\equiv A(x)(mod x^{\lceil\frac{n}{2}\rceil})\\ B(x)'^2&\equiv A(x)(mod x^{\lceil\frac{n}{2}\rceil})\\ B(x)^2-B'(x)^2&\equiv 0(mod x^{\lceil\frac{n}{2}\rceil})\\ B(x)^4+B'(x)^4-2B(x)^2B'(x)^2&\equiv 0(mod x^{\frac{n}{2}})\\ (B(x)^2+B'(x)^2)^2&\equiv (2B(x)B'(x))^2(mod x^{\frac{n}{2}})\\ B(x)^2+B'(x)^2&\equiv 2B(x)B'(x) (mod x^{\frac{n}{2}})\\ A(x)+B'(x)^2&\equiv 2B(x)B'(x)(mod x^{\frac{n}{2}})\\ B(x)&\equiv \frac{A(x)+B'(x)^2}{2B'(x)}(mod x^{\frac{n}{2}})\\ \end{aligned}\]
边界：\(n=1\longrightarrow (mod x)\)，则仅剩常数项，\(B_0=\sqrt{A_0}\)

至此，我们得到的最终的式子，仅需求逆就能实现开方了

code

细节：由于期间在求逆与开方间反复调换，要注意模的次数

#include<bits/stdc++.h>
typedef long long LL;
const LL maxn=1e6+9,mod=998244353,g=3;
inline LL Read(){
    LL x(0),f(1);char c=getchar();
    while(c<'0' || c>'9'){
        if(c=='-') f=-1; c=getchar();
    }
    while(c>='0' && c<='9'){
        x=(x<<3)+(x<<1)+c-'0'; c=getchar();
    }
    return x*f;
}
inline LL Pow(LL base,LL b){
    LL ret(1);
    while(b){
        if(b&1) ret=ret*base%mod; base=base*base%mod; b>>=1;
    }return ret;
}
LL r[maxn];
inline void NTT(LL *a,LL n,LL type){
    for(LL i=0;i<n;++i) if(i<r[i]) std::swap(a[i],a[r[i]]);
    for(LL mid=1;mid<n;mid<<=1){
        LL wn(Pow(g,(mod-1)/(mid<<1)));
        if(type==-1) wn=Pow(wn,mod-2);
        for(LL R=mid<<1,j=0;j<n;j+=R)
            for(LL k=0,w=1;k<mid;++k,w=w*wn%mod){
                LL x(a[j+k]),y(w*a[j+mid+k]%mod);
                a[j+k]=(x+y)%mod; a[j+mid+k]=(x-y+mod)%mod;
            }
    }
    if(type==-1){
        LL ty(Pow(n,mod-2));
        for(LL i=0;i<n;++i) a[i]=a[i]*ty%mod;
    }
}
inline LL Fir(LL n){
    LL limit(1),len(0);
    while(limit<n){
        limit<<=1; ++len;
    }
    for(LL i=0;i<limit;++i) r[i]=(r[i>>1]>>1)|((i&1)<<len-1);
    return limit;
}
LL F[maxn];
void Solve_inv(LL deg,LL *A,LL *B){
    if(deg==1){
        B[0]=Pow(A[0],mod-2); return;
    }
    Solve_inv(deg+1>>1,A,B);
    for(LL i=0;i<deg;++i) F[i]=A[i];
    LL limit(Fir(deg<<1));
    for(LL i=deg;i<limit;++i) F[i]=0;
    NTT(F,limit,1); NTT(B,limit,1);
    for(LL i=0;i<limit;++i)
        B[i]=(2ll-B[i]*F[i]%mod+mod)%mod*B[i]%mod;
    NTT(B,limit,-1);
    for(LL i=deg;i<limit;++i) B[i]=0;
}
LL _B[maxn],C[maxn],T[maxn];
void Solve_kf(LL deg,LL *A,LL *B){
    if(deg==1){
        B[0]=1ll; return;
    }
    Solve_kf(deg+1>>1,A,B);
    LL n(deg+1>>1);
    for(LL i=0;i<n;++i) _B[i]=2ll*B[i]%mod;
    LL limit(Fir(deg<<1));
    for(LL i=0;i<limit;++i) C[i]=0;
    Solve_inv(deg,_B,C);
    for(LL i=0;i<deg;++i) T[i]=A[i];
    for(LL i=deg;i<limit;++i) T[i]=0;
    
    NTT(B,limit,1); NTT(T,limit,1); NTT(C,limit,1);
    for(LL i=0;i<limit;++i) B[i]=(B[i]*B[i]%mod+T[i])%mod*C[i]%mod;
    NTT(B,limit,-1);
    for(LL i=deg;i<limit;++i) B[i]=0;
}
LL n;
LL a[maxn],b[maxn];
int main(){
    n=Read();
    for(LL i=0;i<n;++i) a[i]=Read();
    Solve_kf(n,a,b);
    for(LL i=0;i<n;++i) printf("%lld ",b[i]);
    return 0;
}