Day 075-Q2 — FPS sqrt（Newton法 + NTT による形式的冪級数平方根）

問題

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

$$g(x)^2 \equiv f(x) \pmod{x^N}$$

を満たす $g(x)$ を $\pmod{x^N}$ で求めよ。$a_0$ は $p$ 上で平方数であることが保証される。

入出力例

4
1 2 3 4

1 1 1 3

$g = 1 + x + x^2 + 3x^3$ として $(1+x+x^2+3x^3)^2 \equiv 1+2x+3x^2+4x^3 \pmod{x^4}$ を確認できる。

概念図: Newton 反復による倍精度アップ

ヒント

def fps_sqrt(f, n, MOD=998244353):
    inv2 = (MOD + 1) // 2
    g = [tonelli_shanks(f[0], MOD)]
    k = 1
    while k < n:
        k2 = min(2 * k, n)
        g_inv = fps_inv(g, k2, MOD)
        fg_inv = poly_mul(f[:k2], g_inv)[:k2]
        g = [(( g[i] if i < len(g) else 0) + fg_inv[i]) * inv2 % MOD
             for i in range(k2)]
        k = k2
    return g[:n]

模範解答

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

def ntt(a, invert=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:
        w = pow(3, (MOD - 1) // length, MOD)
        if invert: w = pow(w, MOD - 2, MOD)
        for i in range(0, n, length):
            wn = 1
            for k in range(length // 2):
                u = a[i + k]; v = 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 invert:
        inv_n = pow(n, MOD - 2, MOD)
        for i in range(n): a[i] = a[i] * inv_n % MOD

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

def fps_inv(f, n):
    g = [pow(f[0], MOD - 2, MOD)]; k = 1
    while k < n:
        k2 = min(2*k, n)
        fg = poly_mul(f[:k2], poly_mul(g, g))[:k2]
        g2 = [(2*g[i] if i < len(g) else 0) % MOD for i in range(k2)]
        g = [(g2[i] - fg[i]) % MOD for i in range(k2)]; k = k2
    return g[:n]

def tonelli_shanks(n, p):
    if n == 0: return 0
    if p % 4 == 3: return pow(n, (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(n, q, p), pow(n, (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 fps_sqrt(f, n):
    inv2 = (MOD + 1) // 2
    g = [tonelli_shanks(f[0], MOD)]; k = 1
    while k < n:
        k2 = min(2*k, n)
        g_inv = fps_inv(g, k2)
        fg_inv = poly_mul(f[:k2], g_inv)[:k2]
        g = [((g[i] if i < len(g) else 0) + fg_inv[i]) * inv2 % MOD
             for i in range(k2)]
        k = k2
    return g[:n]

def solve():
    N = int(input())
    a = list(map(int, input().split()))
    print(*fps_sqrt(a, N))

solve()

Step-by-Step 解説

Step 1: Tonelli-Shanks による $\sqrt{a_0} \bmod p$

$p \equiv 3 \pmod{4}$ のとき $\sqrt{n} = n^{(p+1)/4} \bmod p$ で直接計算できる。$p = 998244353$ は $p \equiv 1 \pmod{4}$ なので Tonelli-Shanks の完全版が必要。$p-1 = 2^s \cdot q$ と分解し、非平方数 $z$ を見つけて反復で解を絞り込む。

Step 2: FPS 逆元の Newton 反復

$g^{-1}_{2k} \equiv 2g^{-1}_k - f \cdot (g^{-1}_k)^2 \pmod{x^{2k}}$。これは $h(g) = g^{-1} - f = 0$ の Newton 反復。各ステップで NTT を使い、$O(N \log N)$ 総計算量。

Step 3: FPS sqrt の Newton 反復

$g' = \frac{g + f/g}{2} = \frac{1}{2}\left(g + f \cdot g^{-1}\right)$。先に $g^{-1} \bmod x^{2k}$ を求め、$f \cdot g^{-1}$ を NTT で計算し、$(g + f \cdot g^{-1}) / 2$ が次の $g$ となる。

Step 4: 全体計算量

各ステップで NTT $O(k \log k)$。$k = 1, 2, 4, \ldots, N$ と倍増するので総和は $O(N \log N)$。

よくあるミス

次のステップ

• 発展問題: FPS exp（$e^{f(x)} \bmod x^N$）を同様の Newton 法で実装する（$g' \leftarrow g(1 - \log g + f)$）

自己評価