概述

给定一个 $n - 1$ 次多项式 $A(x)$，你需要求出一个 $n - 1$ 次多项式 $B(x)$ 满足 $B ^ 2(x) \equiv A(x)\pmod{x ^ n}$。

思路

我们还是利用「算法笔记」多项式求逆 的倍增方法。假设我们已经知道了模 $x ^ {\left\lceil\frac{n}{2}\right\rceil}$ 意义下的多项式 $B'(x)$，即有：

$$ B' ^ 2(x) \equiv A(x)\pmod{x ^ {\left\lceil\frac{n}{2}\right\rceil}} $$

又因为：

$$ B ^ 2(x) \equiv A(x)\pmod{x ^ {\left\lceil\frac{n}{2}\right\rceil}} $$

两式相减得到：

$$ B ^ 2(x) - B' ^ 2(x) \equiv 0\pmod{x ^ {\left\lceil\frac{n}{2}\right\rceil}} $$

用平方差公式得到：

$$ [B(x) - B'(x)][B(x) + B'(x)] \equiv 0\pmod{x ^ {\left\lceil\frac{n}{2}\right\rceil}} $$

发现这个式子中 $B(x)$ 有两个解，我们取其中一个解 $B'(x)$ 得到（一般而言取的都是正根，而我们定义常数项较小的根为正根）：

$$ B(x) - B'(x) \equiv 0\pmod{x ^ {\left\lceil\frac{n}{2}\right\rceil}} $$

两边平方后，使用求逆中的推导过程，将其转化为模 $x ^ n$ 意义下：

$$ B ^ 2(x) + B' ^ 2(x) - 2B(x)B'(x) \equiv 0\pmod{x ^ n} $$

由于 $B ^ 2(x) \equiv A(x)\pmod{x ^ n}$，因此式子化为：

$$ A(x) + B' ^ 2(x) - 2B(x)B'(x) \equiv 0\pmod{x ^ n} $$

移项得：

$$ B(x) \equiv \frac{A(x) + B' ^ 2(x)}{2B'(x)} \pmod{x ^ n} $$

递归的边界条件为 $n = 1$ 时 $B(x)$ 为 $A(x)$ 的常数项在模 $998244353$ 的二次剩余，可以通过 $\text{BSGS}$ 在 $\mathcal O(\sqrt{p})$ 的时间内暴力实现。

直接套上一个多项式求逆后迭代实现即可。

时间复杂度：$\mathcal O(n \log n)$。

代码

Vec sqrt(Vec A) {
    int n = A.size(), N = extend(n);
    A.resize(N);
    Vec R(N, 0);
    R[0] = degree(A[0], 2, MOD);
    int i2 = inv(2);
    for (int l = 2; l <= N; l <<= 1) {
        Vec P(l), Q(l);
        std::copy(A.begin(), A.begin() + l, P.begin());
        std::copy(R.begin(), R.begin() + l, Q.begin());
        Vec I = ~Q;
        int L = l << 1;
        P.resize(L), DFT(P);
        Q.resize(L), DFT(Q);
        I.resize(L), DFT(I);
        for (int i = 0; i < L; i++) {
            Q[i] = 1LL * Q[i] * Q[i] % MOD;
            P[i] = 1LL * (P[i] + Q[i]) * i2 % MOD * I[i] % MOD;
        }
        IDFT(P), P.resize(l);
        std::copy(P.begin(), P.begin() + l, R.begin());
    }
    R.resize(n);
    return R;
}