Luogu5205 【模板】多项式开根

https://www.luogu.com.cn/problem/P5205

分治/多项式求逆

\[设G(n)^2 \equiv A(n) (\mod x^{\lceil \frac{n}{2} \rceil})\\ B(n)^2 \equiv A(n) (\mod x^{\lceil \frac{n}{2} \rceil})\\ \therefore G(n) \equiv B(n) (\mod x^{\lceil \frac{n}{2} \rceil})\\ (G(n)-B(n))^2 \equiv 0(\mod x^n)\\ G(n)^2-2G(n)B(n)+B(n)^2 \equiv 0(\mod x^n)\\ G(n)^2-2G(n)B(n)+A(n) \equiv 0(\mod x^n)\\ B(n) \equiv \frac{A(n)+G(n)^2}{2G(n)} (\mod x^n) \]

注意：

\(1.\)代码中的\(lg2\)数组如果要用，一定要开够，否则很容易错

\(C++ Code:\)

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<cstring>
#define N 400005
#define p 998244353
#define ll long long
using namespace std;
int n,k,S;
int lg2[N],rev[N];
ll inv2,a[N],b[N],c[N],d[N],e[N],h[N],A[N],B[N],G[2][24];
ll ksm(ll x,ll y)
{
    ll ans=1;
    while (y)
    {
        if (y & 1)
            ans=ans*x%p;
        x=x*x%p;
        y >>=1;
    }
    return ans;
}
#define inv(x) (ksm(x,p-2))
void NTT(ll *a,int t)
{
    for (int i=0;i<S;i++)
        if (i<rev[i])
            swap(a[i],a[rev[i]]);
    for (int mid=1,o=1;mid<S;mid <<=1,o++)
    {
        ll gn=G[t][o];
        for (int j=0;j<S;j+=(mid << 1))
        {
            ll g=1;
            for (int k=0;k<mid;k++,g=g*gn%p)
            {
                ll x=a[j+k],y=g*a[j+k+mid]%p;
                a[j+k]=(x+y)%p;
                a[j+k+mid]=(x-y)%p;
            }
        }
    }
}
void Pre()
{
    inv2=inv(2);
    G[0][23]=ksm(3,(p-1)/(1 << 23));
    G[1][23]=inv(G[0][23]);
    for (int i=22;i>=1;i--)
    {
        G[0][i]=G[0][i+1]*G[0][i+1]%p;
        G[1][i]=G[1][i+1]*G[1][i+1]%p;
    }
}
void Qn(int r)
{
    if (!r)
    {
        h[r]=inv(b[r]);
        return;
    }
    int mid=r >> 1;
    Qn(mid);
    int l=lg2[r+1]+1;
    int s=1 << l;
    for (int i=0;i<s;i++)
        rev[i]=(rev[i >> 1] >> 1) | ((i & 1) << (l-1));
    for (int i=0;i<=mid;i++)
        c[i]=h[i];
    for (int i=mid+1;i<s;i++)
        c[i]=0;
    for (int i=0;i<=r;i++)
        d[i]=b[i];
    for (int i=r+1;i<s;i++)
        d[i]=0;
    S=s;
    NTT(c,0);
    NTT(d,0);
    for (int i=0;i<s;i++)
        c[i]=c[i]*c[i]%p*d[i]%p;
    NTT(c,1);
    ll t=inv(s);
    for (int i=0;i<s;i++)
        c[i]=c[i]*t%p;
    for (int i=0;i<=r;i++)
        h[i]=(h[i]*2-c[i])%p;
}
void FZ(int r)
{
    if (!r)
    {
        B[r]=1;
        return;
    }
    int mid=r >> 1;
    FZ(mid);
    int l=lg2[r+1]+1;
    int s=1 << l;
    for (int i=0;i<=mid;i++)
        b[i]=B[i]%p;
    for (int i=mid+1;i<s;i++)
        b[i]=0;
    for (int i=0;i<s;i++)
        h[i]=0;
    Qn((s >> 1)-1);
    for (int i=0;i<=mid;i++)
        a[i]=B[i];
    for (int i=mid+1;i<s;i++)
        a[i]=0;
    for (int i=0;i<=r;i++)
        e[i]=A[i];
    for (int i=r+1;i<s;i++)
        e[i]=0;
    S=s;
    for (int i=0;i<s;i++)
        rev[i]=(rev[i >> 1] >> 1) | ((i & 1) << (l-1));
    NTT(a,0);
    NTT(h,0);
    NTT(e,0);
    for (int i=0;i<s;i++)
        a[i]=(a[i]*a[i]%p+e[i])%p*h[i]%p;
    NTT(a,1);
    ll t=inv(s);
    for (int i=0;i<s;i++)
        a[i]=a[i]*t%p*inv2%p;
    for (int i=0;i<=r;i++)
        B[i]=a[i];
}
int main()
{
    Pre();
    scanf("%d",&n);
    for (int i=0;i<n;i++)
        scanf("%lld",&A[i]);
    lg2[0]=0;
    int l=0;
    for (int i=1;i<=n;i++)
    {
        if ((1 << l)<i)
            l++;
        lg2[i]=l;
    }
    for (int i=n+1;i<=(1 << lg2[n]);i++)
        lg2[i]=lg2[n];
    for (int i=(1 << lg2[n])+1;i<=(1 << (lg2[n]+1));i++)
        lg2[i]=lg2[n]+1;
    FZ((1 << lg2[n])-1);
    for (int i=0;i<n;i++)
        B[i]=(B[i]%p+p)%p;
    for (int i=0;i<n;i++)
        printf("%lld ",B[i]);
    putchar('\n');
    return 0;
}