洛谷P5205 【模板】多项式开根题解

题意：

给定一个 \(n-1\) 次多项式 \(A(x)\)，求一个在 \({} \bmod x^n\) 意义下的多项式 \(B(x)\)，使得 \(B^2(x) \equiv A(x) \pmod{x^n}\)。若有多解，请取零次项系数较小的作为答案。

多项式的系数在模 \(998244353\) 意义下进行运算，本题保证 \(a_0 = 1\)。

输入格式：

第一行一个正整数 \(n\)。

接下来 \(n\) 个整数，依次表示多项式的系数 \(a_0, a_1, \dots, a_{n-1}\)。

输出格式：

输出 \(n\) 个整数，表示答案多项式的系数 \(b_0, b_1, \dots, b_{n-1}\)。

数据范围：

对于 \(100 \%\) 的数据：\(1 \le n \leq 10^5\)，\(0 \le a_i < 998244353\)。

设答案为 \(g(x)\) 。因为 \([x^0]g(x)\) 不为 \(0\) ，易知 \[ [x^0]f(x) = \sqrt{[x^0]g(x)} \] 若 \([x^0]g(x)\) 没有平方根，则多项式 \(g(x)\) 没有平方根

另外，\(\left[x^0\right] g(x)\) 可能有多个平方根，因此选取不同的根会求出不同的 \(f(x)\) 。

假设现在已经求出了 \(g(x)\) 在模 \(x^{\left\lceil\frac{n}{2}\right\rceil}\) 意义下的平方根 \(f_0(x)\) ，则有 \[ \begin{aligned} f_0^2(x) & \equiv g(x) \pmod{x^{\left\lceil\frac{n}{2}\right\rceil}} \\[6pt]f_0^2(x)-g(x) & \equiv 0 \pmod{x^{\left\lceil\frac{n}{2}\right\rceil}} \\[6pt]\left(f_0^2(x)-g(x)\right)^2 & \equiv 0 \pmod{x^{n}} \\[6pt]\left(f_0^2(x)+g(x)\right)^2 & \equiv 4 f_0^2(x) g(x) \pmod{x^{n}} \\[6pt]\left(\frac{f_0^2(x)+g(x)}{2 f_0(x)}\right)^2 & \equiv g(x) \pmod{x^{n}} \\[6pt]\frac{f_0^2(x)+g(x)}{2 f_0(x)} & \equiv f(x) \pmod{x^{n}} \\[6pt]2^{-1} f_0(x)+2^{-1} f_0^{-1}(x) g(x) & \equiv f(x) \end{aligned} \] 于是倍增计算即可。

时间复杂度 \(T(n)=T\left(\frac{n}{2}\right)+\mathcal{O}(n \log n)=\mathcal{O}(n \log n)\)

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
void up(int &x,int y) { x < y ? x = y : 0; }
void down(int &x,int y) { x > y ? x = y : 0; }
#define N ((int)(2097152 + 15))

const int P = 998244353;
const int G = 3, Gi = 332748118;
int a[N], b[N], c[N], d[N], r[N];
int qpow(int a, int b)
{
    int r = 1;
    while(b)
    {
        if(b & 1) r = r * a % P;
        b >>= 1; a = a * a % P;
    }
    return r;
}
void NTT(int *A, int limit, int type)
{
    for(int i = 0; i < limit; i++) if(i < r[i]) swap(A[i], A[r[i]]);
    for(int mid = 1; mid < limit; mid *= 2)
    {
        int Wn = qpow((type == 1 ? G : Gi), (P - 1) / (mid * 2));
        for(int j = 0; j < limit; j += (mid * 2))
        {
            int w = 1;
            for(int k = 0; k < mid; k++, w = (w * Wn) % P)
            {
                int x = A[j + k], y = w * A[j + k + mid] % P;
                A[j + k] = (x + y) % P; A[j + k + mid] = (x - y + P) % P;
            }
        }
    }
    if(type == -1)
    {
        int Inv = qpow(limit, P - 2);
        for(int i = 0; i < limit; i++) A[i] = A[i] * Inv % P;
    }
}
void Inv(int n, int *a, int *b)
{
    if(n == 1) { b[0] = qpow(a[0], P - 2); return; }
    Inv((n + 1) / 2, a, b);
    int l = 0, limit = 1;
    for(; limit <= (n - 1) * 2; l++, limit *= 2);
    for(int i = 0; i < limit; i++) r[i] = (r[i >> 1] >> 1) | ((i & 1) << (l - 1));

    for(int i = 0; i < n; i++) c[i] = a[i];
    for(int i = n; i < limit; i++) c[i] = 0;

    NTT(c, limit, 1); NTT(b, limit, 1);
    for(int i = 0; i < limit; i++)
        b[i] = (2 - c[i] * b[i] % P + P) % P * b[i] % P;
    NTT(b, limit, -1);
    for(int i = n; i < limit; i++) b[i] = 0;
}
void Sqrt(int n, int *a, int *b)
{
    if(n == 1) { b[0] = 1; return; }
    Sqrt((n + 1) / 2, a, b);
    int l = 0, limit = 1;
    for(; limit <= (n - 1) * 2; l++, limit *= 2);
    for(int i = 0; i < limit; i++) r[i] = (r[i >> 1] >> 1 | ((i & 1) << (l - 1)));
    for(int i = 0; i < limit; i++) d[i] = 0;
    Inv(limit / 2, b, d);
    for(int i = 0; i < n; i++) c[i] = a[i];
    for(int i = n; i < limit; i++) c[i] = 0;

    NTT(c, limit, 1); NTT(d, limit, 1);
    for(int i = 0; i < limit; i++) c[i] = c[i] * d[i] % P;
    NTT(c, limit, -1);
     static int inv2 = qpow(2, P - 2);
    for(int i = 0; i < n; i++) b[i] = (b[i] + c[i]) % P * inv2 % P;
    for(int i = n; i < limit; i++) b[i] = 0;
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    int n; cin >> n;
    for(int i = 0, x; i < n; i++) cin >> x, a[i] = (x + P) % P;
    Sqrt(n, a, b);
    for(int i = 0; i < n; i++) cout << b[i] << " \n"[i == n - 1];
    return 0;
}