Day 134-Q3 — Toom-Cook法（Toom-3）による多項式乗算

2026-08-26 赤色 Master / Phase 8+ ★★★★★★★★★ 3分割5点評価による部分積削減 O(N^log_3 5)

問題

2つの多項式 $A(x)=\sum_{i=0}^{N-1}a_ix^i$、$B(x)=\sum_{i=0}^{M-1}b_ix^i$（係数はmod $998244353$）が与えられる。積 $C(x)=A(x)B(x)=\sum_{i=0}^{N+M-2}c_ix^i$ の係数をすべてmod $998244353$で求めよ。

入力形式

N M
a_0 a_1 ... a_{N-1}
b_0 b_1 ... b_{M-1}

制約

$1 \le N,M \le 20000$
$0 \le a_i,b_i < 998244353$
法は常に$998244353$（NTT-friendly素数）

入出力例

入力例1

3 3
1 2 3
4 5 6

出力例1

4 13 28 27 18

概念図

Toom-3の3段構成: 分割 → 5点で評価・積 → 補間 f0 | f1 | f2 A,Bを3ブロックに分割 評価点 y=0,1,-1,2,∞ f(y)・g(y) を5組計算 (9回→5回の部分積に削減) 各評価は再帰 or 愚直乗算 補間 c0..c4を 連立方程式で復元 (inv2,inv6使用) C(x) = c0 + c1·x^k + c2·x^2k + c3·x^3k + c4·x^4k（シフトして加算） 再帰: T(N) = 5T(N/3) + O(N) ⇒ O(N^(log_3 5)) ≈ O(N^1.465)

ヒント（段階的開示）

ヒント1（方向性）

NTTを使えば$O(N\log N)$で解けるが、本問はあえてNTTを使わずKaratsuba法の3分割版であるToom-Cook法（Toom-3）を実装する。$A,B$を3ブロックに分割し、通常なら$3\times3=9$回必要なブロック積を5回に減らすのが核心。

ヒント2（アプローチ）

$f(y)=f_0+f_1y+f_2y^2$（$y=x^k$）とみなし、5点$y=0,1,-1,2,\infty$での値を計算する（$g$も同様）。各点で$f(y)\cdot g(y)$を計算（ブロック同士の多項式乗算、再帰または愚直）。得られた5点の積から次数4の多項式の係数を補間で復元する。

ヒント3（誘導） 
評価: f(0)=f0, f(1)=f0+f1+f2, f(-1)=f0-f1+f2, f(2)=f0+2f1+4f2, f(inf)=f2
積:   r0=f(0)g(0), r1=f(1)g(1), rm1=f(-1)g(-1), r2=f(2)g(2), rinf=f(inf)g(inf)
補間: c0=r0
      c4=rinf
      c2=(r1+rm1-2c0-2c4)*inv(2)
      c3=(r2-c0-16c4-4c2-r1+rm1)*inv(6)
      c1=(r1-rm1)*inv(2)-c3

ブロックサイズが一定以下になったら再帰を打ち切り愚直な$O(nm)$乗算に切り替える。

模範解答 (Python)

import sys

MOD = 998244353
INV2 = pow(2, MOD - 2, MOD)
INV6 = pow(6, MOD - 2, MOD)
BASE_CASE = 24  # このサイズ以下は愚直乗算に切り替える


def naive_mul(a, b):
    n, m = len(a), len(b)
    res = [0] * (n + m - 1)
    for i in range(n):
        ai = a[i]
        if ai == 0:
            continue
        for j in range(m):
            res[i + j] = (res[i + j] + ai * b[j]) % MOD
    return res


def padd(a, b):
    n = max(len(a), len(b))
    res = [0] * n
    for i, v in enumerate(a):
        res[i] = (res[i] + v) % MOD
    for i, v in enumerate(b):
        res[i] = (res[i] + v) % MOD
    return res


def psub(a, b):
    n = max(len(a), len(b))
    res = [0] * n
    for i, v in enumerate(a):
        res[i] = (res[i] + v) % MOD
    for i, v in enumerate(b):
        res[i] = (res[i] - v) % MOD
    return res


def pscale(a, c):
    return [(x * c) % MOD for x in a]


def toom3_mul(a, b):
    n = max(len(a), len(b))
    if n <= BASE_CASE:
        return naive_mul(a, b)

    k = (n + 2) // 3

    def split(p):
        p = p + [0] * max(0, 3 * k - len(p))
        return p[0:k], p[k:2 * k], p[2 * k:3 * k]

    f0, f1, f2 = split(a[:])
    g0, g1, g2 = split(b[:])

    fe0, ge0 = f0, g0
    fe1 = padd(padd(f0, f1), f2)
    ge1 = padd(padd(g0, g1), g2)
    fem1 = psub(padd(f0, f2), f1)
    gem1 = psub(padd(g0, g2), g1)
    fe2 = padd(padd(f0, pscale(f1, 2)), pscale(f2, 4))
    ge2 = padd(padd(g0, pscale(g1, 2)), pscale(g2, 4))
    feinf, geinf = f2, g2

    r0 = toom3_mul(fe0, ge0)
    r1 = toom3_mul(fe1, ge1)
    rm1 = toom3_mul(fem1, gem1)
    r2 = toom3_mul(fe2, ge2)
    rinf = toom3_mul(feinf, geinf)

    c0 = r0
    c4 = rinf
    c2 = pscale(psub(padd(r1, rm1), padd(pscale(c0, 2), pscale(c4, 2))), INV2)
    tmp = psub(r2, c0)
    tmp = psub(tmp, pscale(c4, 16))
    tmp = psub(tmp, pscale(c2, 4))
    tmp = psub(tmp, r1)
    tmp = padd(tmp, rm1)
    c3 = pscale(tmp, INV6)
    c1 = psub(pscale(psub(r1, rm1), INV2), c3)

    outlen = len(a) + len(b) - 1
    res = [0] * max(outlen, 4 * k + len(c4))

    def add_shifted(c, shift):
        for i, v in enumerate(c):
            if v:
                res[shift + i] = (res[shift + i] + v) % MOD

    add_shifted(c0, 0)
    add_shifted(c1, k)
    add_shifted(c2, 2 * k)
    add_shifted(c3, 3 * k)
    add_shifted(c4, 4 * k)
    return res[:outlen]


def main():
    sys.setrecursionlimit(10000)
    data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(data[idx]); idx += 1
    M = int(data[idx]); idx += 1
    a = [int(data[idx + i]) for i in range(N)]; idx += N
    b = [int(data[idx + i]) for i in range(M)]; idx += M
    c = toom3_mul(a, b)
    print(*c)


main()

Step-by-Step 解説

1なぜ9回ではなく5回か
ブロックを係数とする2次式$f,g$の積（4次式）は5点での評価値から一意に復元できる、という事実を利用する。
2評価点の選び方
$0,\pm1,\infty$に加え計算が軽い$2$を追加。$\infty$での評価は最高次係数を取り出すだけで実質タダ。
3補間の代数
5連立方程式から$c_0\to c_4\to c_2\to c_3\to c_1$の順に解く。$c_2$に2の逆元、$c_3$に6の逆元が必要。
4再帰と基底ケース
ブロックが十分小さくなったら愚直乗算に切り替え、定数倍のオーバーヘッドを抑える。
5計算量比較
Toom-3は$O(N^{1.465})$、Karatsubaは$O(N^{1.585})$、NTTは$O(N\log N)$。漸近的にはNTTに劣るが定数倍は小さい。

よくあるミス

ミス原因正しい書き方
補間の計算順序を間違える連立方程式の依存関係を見落とす必ず$c_0\to c_4\to c_2\to c_3\to c_1$の順に計算
modでの減算が負のまま次の計算に渡る複数回の減算を1式にまとめて意図せず負が残る各演算関数内で必ず% MODを通す
ブロック分割時の0埋め不足配列長が$3k$の倍数でないとインデックスがずれるsplit関数で必ず$3k$になるよう0埋めしてから3等分
基底ケースが小さすぎ再帰オーバーヘッドが支配的「サイズ1まで再帰」が理論的に正しいと思い込む基底ケースを十分大きく取り、そこから先は愚直乗算に切り替える
sys.setrecursionlimitを設定し忘れる再帰深さは$O(\log_3 N)$で浅いが安全マージンが薄い念のため明示的に引き上げておく

次のステップ

  • 発展: 本問のNTT版を実装し、$N=20000$での実行時間を比較せよ。
  • 発展2: 分割数を3から4,5と増やしたToom-4,Toom-5の一般化を調べ、漸近改善と実装複雑化のトレードオフを整理せよ。
  • 次回予告: 次回のテーマは今後のセッションで選定

自己評価

自分の回答

気づき・メモ