Day 053-Q4 | 多項式平方根・逆元（FPS sqrt/inv over NTT + Newton法）

問題

素数 $p = 998244353$ 上の形式的冪級数 $f(x) = \sum_{i=0}^{N-1} a_i x^i$ が与えられる。以下を求めよ。

$f(x)$ の逆元 $g(x)$: $f \cdot g \equiv 1 \pmod{x^N}$
$f(x)$ の平方根 $h(x)$: $h^2 \equiv f \pmod{x^N}$（$a_0$ は完全平方数を保証）

制約

パラメータ	範囲
$N$	$1 \le N \le 2 \times 10^5$
$a_i$	$0 \le a_i < 998244353$
$a_0$	$\ne 0$（逆元の存在を保証）
$a_0$	$\mathbb{F}_p$ 上の完全平方数
時間制限	2秒

入出力例

入力例 1

4
1 2 3 4

出力例 1

逆元: 1 998244351 1 998244347
平方根: 1 1 1 3

$f = 1+2x+3x^2+4x^3$。逆元 $g$: $fg \equiv 1 \pmod{x^4}$。平方根 $h$: $h^2 \equiv f \pmod{x^4}$。

概念図: Newton法による精度倍増

ヒント（段階的開示）

ヒント1: 方向性

FPS の逆元・平方根はどちらも Newton法（精度倍増）で計算できる。「精度 $n$」の近似から「精度 $2n$」の近似を $O(n \log n)$ で得る。$\log N$ 回の反復で $O(N \log N)$ で収束する。

ヒント2: アプローチ

逆元: $\phi(g) = g^{-1} - f = 0$ の Newton反復: $g_{k+1} = g_k(2 - f g_k) \pmod{x^{2n}}$
平方根: $\phi(h) = h^2 - f = 0$ の Newton反復: $h_{k+1} = (h_k + f h_k^{-1}) / 2 \pmod{x^{2n}}$
初期値: $g_0 = a_0^{-1}$, $h_0 = \sqrt{a_0}$ を $\mathbb{F}_p$ 上で計算（Tonelli-Shanks）。

ヒント3: コード骨格

def poly_inv(f, n):
    g = [pow(f[0], MOD-2, MOD)]
    size = 1
    while size < n:
        size <<= 1
        # g = g * (2 - f*g) mod x^size
        fg = poly_mul(f[:size], g)[:size]
        neg_fg = [(MOD - v) % MOD for v in fg]
        neg_fg[0] = (neg_fg[0] + 2) % MOD
        g = poly_mul(g, neg_fg)[:size]
    return g[:n]

def poly_sqrt(f, n):
    s0 = tonelli_shanks(f[0])  # sqrt(a_0) in F_p
    h = [s0]
    inv2 = (MOD + 1) // 2
    size = 1
    while size < n:
        size <<= 1
        hi = poly_inv(h, size)
        fhi = poly_mul(f[:size], hi)[:size]
        h = [((hj + fj) * inv2) % MOD
             for hj, fj in zip(h + [0]*(size-len(h)), fhi)][:size]
    return h[:n]

模範解答 (Python)

import sys
input = sys.stdin.readline
MOD = 998244353
PRIM_ROOT = 3

def ntt(a, inv=False):
    n = len(a)
    j = 0
    for i in range(1, n):
        bit = n >> 1
        while j & bit: j ^= bit; bit >>= 1
        j ^= bit
        if i < j: a[i], a[j] = a[j], a[i]
    length = 2
    while length <= n:
        exp = (MOD-1) // length
        if inv: exp = (MOD-1) - exp
        w = pow(PRIM_ROOT, exp, MOD)
        for i in range(0, n, length):
            wn = 1
            for k in range(length // 2):
                u, v = a[i+k], a[i+k+length//2] * wn % MOD
                a[i+k] = (u + v) % MOD
                a[i+k+length//2] = (u - v) % MOD
                wn = wn * w % MOD
        length <<= 1
    if inv:
        ni = pow(n, MOD-2, MOD)
        for i in range(n): a[i] = a[i] * ni % MOD

def poly_mul(a, b, lim=None):
    n = 1
    need = len(a) + len(b) - 1
    while n < need: n <<= 1
    fa = a + [0] * (n - len(a))
    fb = b + [0] * (n - len(b))
    ntt(fa); ntt(fb)
    fc = [fa[i] * fb[i] % MOD for i in range(n)]
    ntt(fc, inv=True)
    return fc[:lim] if lim else fc

def poly_inv(f, n):
    g = [pow(f[0], MOD-2, MOD)]
    size = 1
    while size < n:
        size <<= 1
        fg = poly_mul(f[:size], g, size)
        neg = [(MOD - v) % MOD for v in fg]
        neg[0] = (neg[0] + 2) % MOD
        g = poly_mul(g, neg, size)
    return g[:n]

def tonelli_shanks(a, p=MOD):
    if a == 0: return 0
    if p % 4 == 3: return pow(a, (p+1)//4, p)
    q, s = p-1, 0
    while q % 2 == 0: q //= 2; s += 1
    z = 2
    while pow(z, (p-1)//2, p) != p-1: z += 1
    m, c, t, r = s, pow(z, q, p), pow(a, q, p), pow(a, (q+1)//2, p)
    while True:
        if t == 1: return r
        i, tmp = 1, t*t % p
        while tmp != 1: tmp = tmp*tmp % p; i += 1
        b = pow(c, 1 << (m-i-1), p)
        m, c, t, r = i, b*b%p, t*b*b%p, r*b%p

def poly_sqrt(f, n):
    s0 = tonelli_shanks(f[0])
    h = [s0]
    inv2 = (MOD + 1) // 2
    size = 1
    while size < n:
        size <<= 1
        hi = poly_inv(h, size)
        fhi = poly_mul(f[:size], hi, size)
        h_ext = h + [0] * (size - len(h))
        h = [((h_ext[i] + fhi[i]) * inv2) % MOD for i in range(size)]
    return h[:n]

def solve():
    N = int(input())
    f = list(map(int, input().split()))
    g = poly_inv(f, N)
    h = poly_sqrt(f, N)
    print("逆元:", *g)
    print("平方根:", *h)

solve()

Step-by-Step 解説

1NTT（数論変換）
FFT の mod $p$ 版。素数 $p = 998244353 = 119 \cdot 2^{23} + 1$ は NTT-friendly（$2^{23}$ 以下の長さで NTT 可能）。原始根 $g = 3$。1の $n$ 乗根 $\omega = g^{(p-1)/n}$。

2逆元の Newton法
$g^2(x)f(x) \equiv g(x) \cdot 2 - g(x)^2 f(x) \pmod{x^{2n}}$。精度が $n$ のとき $2n$ にする反復。毎回多項式乗算 $O(n \log n)$ を 2〜3 回。合計 $O(N \log N)$（等比級数 $\sum_{k=0}^{\log N} O(2^k \log 2^k)$）。

3平方根の Newton法
$\phi(h) = h^2 - f = 0$ の Newton反復。$h_{k+1} = h_k - \phi(h_k)/\phi'(h_k) = (h_k + f/h_k)/2$。$f/h_k$ は $f \cdot h_k^{-1}$。各反復で逆元計算 $O(n \log n)$ + 乗算。

4Tonelli-Shanks（定数項の平方根）
$\mathbb{F}_p$ 上で $a_0$ の平方根を求める。$p \equiv 3 \pmod 4$ なら $\sqrt{a} = a^{(p+1)/4}$。一般（$p \equiv 1 \pmod 4$）には Tonelli-Shanks $O(\log^2 p)$。

計算量

NTT: $O(N \log N)$ 時間（1回）
poly_mul: $O(N \log N)$ — NTT × 3回
poly_inv: $O(N \log N)$ — Newton法（等比級数収束）
poly_sqrt: $O(N \log N)$ — poly_inv × $O(\log N)$ 回
全体: $O(N \log N)$

よくあるミス

ミス	原因	正しい書き方
`2 - fg` の定数項が 0 になる	初期値 `g[0] = a₀⁻¹` が間違い	`pow(f[0], MOD-2, MOD)`
NTT の長さが $2^k$ でない	poly_mul で長さが 2 冪に丸め忘れ	`while n < need: n <<= 1`
`inv2` の計算ミス	$2$ の逆元は $(p+1)/2$（$p$ が奇素数）	確認: `2 * inv2 % MOD == 1`
Newton法の反復回数不足	`size < n` を確認しない	`while size < n: size <<= 1`

次のステップ

発展問題: FPS の log と exp（$\exp(f)$ は $g' = f'g$ の ODE として Newton法で解く）
関連: FPS の pow(f, k) = exp(k * log(f)) — 整数冪も $O(N \log N)$
応用: 多項式行列式・逆行列（Berkowitz algorithm）

自己評価

完全理解おおむね理解要復習