Day 045-Q5 | FPS多項式除算（NTT + Newton法）

問題

$N$ 次多項式 $F(x)$ と $M$ 次多項式 $G(x)$（$N \ge M$, $g_0 \ne 0$）が与えられる。多項式除算 $F(x) = Q(x) \cdot G(x) + R(x)$（$\deg Q = N - M$, $\deg R < M$）における 商 $Q(x)$ と余り $R(x)$ の係数を全て mod $998244353$ で求めよ。

制約

$1 \le M \le N \le 2 \times 10^5$

$0 \le f_i, g_j < 998244353$

$g_0 \ne 0$

時間制限: 2秒

$O(N \log N)$ で解くこと

入出力例

入力例 1

3 1
6 11 6 1
2 1

出力例 1

Q: 3 4 1
R: 0

概念図: 反転多項式を用いた多項式除算

ヒント（段階的開示）

ヒント1: 方向性

通常の多項式除算は $O(N^2)$ ですが、NTT を使えば $O(N \log N)$ に改善できます。鍵となるのは「反転多項式」を使った変換で、除算を乗算に変換します。

ヒント2: アプローチ

$F^{\text{rev}}(x) = x^N F(1/x)$（係数を逆順にした多項式）を定義
$Q^{\text{rev}}(x) \equiv F^{\text{rev}}(x) \cdot (G^{\text{rev}})^{-1}(x) \pmod{x^{N-M+1}}$
$Q$ が求まったら $R = F - Q \cdot G$ を NTT で計算

ヒント3: 実装骨格

def poly_inv(a, n):  # a mod x^n の逆元（Newton法）
    res = [pow(a[0], MOD-2, MOD)]
    length = 1
    while length < n:
        length <<= 1
        # res = 2*res - a*res^2 mod x^length
        r2 = poly_mul_mod(res, res)[:length]
        ar2 = poly_mul_mod(a[:length], r2)[:length]
        res = [(2*res[i] - ar2[i]) % MOD
               if i < len(res) else (-ar2[i]) % MOD
               for i in range(length)]
    return res[:n]

# 反転多項式を使った除算
Fr, Gr = F[::-1], G[::-1]
Gr_inv = poly_inv(Gr, N-M+1)
Qr = poly_mul_mod(Fr, Gr_inv)[:N-M+1]
Q = Qr[::-1]
R = [(F[i] - sum(Q[j]*G[i-j] for ...)) % MOD ...]

模範解答 (Python)

import sys

MOD = 998244353
G_PRIM = 3

def ntt(a, invert):
    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(G_PRIM, (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:
        ni = pow(n, MOD - 2, MOD)
        for i in range(n):
            a[i] = a[i] * ni % MOD

def poly_mul_trunc(a, b, trunc=None):
    if not a or not b:
        return []
    rs = len(a) + len(b) - 1
    n = 1
    while n < rs:
        n <<= 1
    fa = a + [0] * (n - len(a))
    fb = b + [0] * (n - len(b))
    ntt(fa, False); ntt(fb, False)
    for i in range(n):
        fa[i] = fa[i] * fb[i] % MOD
    ntt(fa, True)
    result = fa[:rs]
    if trunc is not None:
        result = result[:trunc]
    return result

def poly_inv(a, n):
    res = [pow(a[0], MOD - 2, MOD)]
    length = 1
    while length < n:
        length <<= 1
        a_trunc = a[:length]
        r2 = poly_mul_trunc(res, res, length)
        ar2 = poly_mul_trunc(a_trunc, r2, length)
        new_res = [0] * length
        for i in range(min(len(res), length)):
            new_res[i] = 2 * res[i] % MOD
        for i in range(len(ar2)):
            new_res[i] = (new_res[i] - ar2[i]) % MOD
        res = new_res
    return res[:n]

def poly_div_rem(F, G):
    n = len(F) - 1
    m = len(G) - 1
    if n < m:
        return [0], F[:]
    deg_Q = n - m
    Fr = F[::-1]
    Gr = G[::-1]
    Gr_inv = poly_inv(Gr, deg_Q + 1)
    Qr = poly_mul_trunc(Fr, Gr_inv, deg_Q + 1)
    Q = Qr[::-1]
    QG = poly_mul_trunc(Q, G)
    R = [(F[i] - (QG[i] if i < len(QG) else 0)) % MOD for i in range(m)]
    return Q, R

def main():
    data = sys.stdin.read().split()
    idx = 0
    N, M = int(data[idx]), int(data[idx+1]); idx += 2
    F = [int(data[idx+i]) for i in range(N+1)]; idx += N+1
    G = [int(data[idx+i]) for i in range(M+1)]; idx += M+1

    Q, R = poly_div_rem(F, G)

    print("Q:", ' '.join(map(str, Q)))
    while R and R[-1] == 0:
        R.pop()
    if not R:
        print("R: 0")
    else:
        print("R:", ' '.join(map(str, R)))

main()

Step-by-Step 解説

1多項式除算の高速化
通常の多項式除算は $O(N^2)$ だが、NTT を使えば $O(N \log N)$ に改善できる。鍵は「反転多項式」を使って除算を乗算に変換すること。

2反転多項式の利用
$F^{\text{rev}}(x) = x^{\deg F} F(1/x)$（係数を逆順）を定義する。 $F = Q \cdot G + R$ の両辺を $x^N$ で整理すると $\pmod{x^{N-M+1}}$ で $R$ の項が消え、 $$Q^{\text{rev}} \equiv F^{\text{rev}} \cdot (G^{\text{rev}})^{-1} \pmod{x^{N-M+1}}$$

3形式的冪級数の逆元（Newton法）
$A(x) \cdot B(x) \equiv 1 \pmod{x^n}$ を Newton 法で求める。倍増アルゴリズム: $B_{2k} = 2B_k - A \cdot B_k^2 \pmod{x^{2k}}$（各ステップ $O(k \log k)$ → 合計 $O(n \log n)$）

4余り R の計算
$Q$ が求まったら $R = F - Q \cdot G$ を NTT で $O(N \log N)$ で計算する。

5NTT（数論変換）
mod $998244353$ ($= 119 \times 2^{23} + 1$) は NTT 素数で、原始根は 3。Cooley-Tukey アルゴリズムで $O(N \log N)$。

計算量

反転: $O(N)$
poly_inv（Newton法）: $O(N \log N)$
poly_mul: $O(N \log N)$
全体: $O(N \log N)$
空間: $O(N)$

よくあるミス

ミス	原因	正しい書き方
poly_inv の終了条件	`>=n` でなく `>n`	`while length < n: length <<= 1`
反転後の配列長を切り詰めない	余分な係数が混入	`Qr = Qr[:deg_Q+1]`
NTT サイズが2の冪でない	NTT が正しく動作しない	`while n < rs: n <<= 1`
mod p での負数処理（C++）	負の剰余が生じる	Python の `%` は常に非負なので問題なし

次のステップ

発展問題: FPS のべき乗 $F(x)^K \pmod{x^N}$（$K \le 10^{18}$）を $O(N \log N \log K)$ で計算する
類題: 多項式 GCD（Euclidean Algorithm for polynomials）
応用: 線形漸化式の $N$ 項目（Bostan-Mori アルゴリズム）

自己評価

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