Day 023-Q4 — 高速フーリエ変換応用 (NTT + 多項式乗算)

2026-05-06 赤色 Master / Phase 8+ ★★★★★★★★★ NTT / Cooley-Tukey

問題

多項式 $f(x) = \sum a_i x^i$(次数 $N-1$)と $g(x) = \sum b_j x^j$(次数 $M-1$)の積 $h(x) = f(x) g(x)$ の係数列を $\bmod 998244353$ で求めよ。

$998244353 = 119 \times 2^{23} + 1$ は NTT 素数。原始根は 3。

制約

$1 \le N, M \le 5 \times 10^5$
$0 \le a_i, b_j < 998244353$

入出力例

入力例 1

3 3
1 2 3
4 5 6

出力例 1

4 13 28 27 18

$(1+2x+3x^2)(4+5x+6x^2) = 4 + 13x + 28x^2 + 27x^3 + 18x^4$

ヒント (段階的開示)

ヒント1: 方向性
NTT は FFT の整数 mod 版。NTT 素数を使えば浮動誤差なしで $O((N+M)\log(N+M))$ の多項式乗算が可能。
ヒント2: アプローチ
①長さを $2^k \ge N+M-1$ にゼロパディング ②NTT で点値に変換 ③点ごとに積 ④INTT で係数に戻す。単位根 $\omega = g^{(p-1)/n} \bmod p$。
ヒント3: 誘導
# Cooley-Tukey バタフライ
length = 2
while length <= n:
    w = pow(G, (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

模範解答 (Python)

import sys
input = sys.stdin.readline

def solve():
    MOD = 998244353
    G_ROOT = 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_ROOT, (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):
        result_len = len(a) + len(b) - 1
        n = 1
        while n < result_len:
            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)
        return fa[:result_len]

    N, M = map(int, input().split())
    a = list(map(int, input().split()))
    b = list(map(int, input().split()))
    print(*poly_mul(a, b))

solve()

Step-by-Step 解説

1NTT の基本
$\omega_n \equiv g^{(p-1)/n} \pmod p$ を使い、FFT を整数体上で実行。$998244353 - 1 = 119 \times 2^{23}$ なので $n \le 2^{23}$ まで OK。
2Cooley-Tukey
DFT を再帰分割し、ビット反転置換でインプレース化。計算量 $O(n \log n)$。
3逆 NTT
$\omega$ を $\omega^{-1}$ に替えて NTT し、最後に $n^{-1}$ を掛ける。
4全体フロー
パディング → NTT($f$), NTT($g$) → 点積 → INTT。$O((N+M)\log(N+M))$。

よくあるミス

ミス原因正しい書き方
NTT 素数でない mod を使う2の冪乗根が存在しない998244353 や 469762049 を使う
長さを 2 の冪にしないNTT の前提条件while n < result_len: n <<= 1
逆変換で $1/n$ 倍を忘れるINTT の正規化a[i] = a[i] * inv_n % MOD

次のステップ

  • 形式的冪級数の逆元・平方根・log・exp
  • 任意 mod での畳み込み(3-NTT)
  • Karatsuba 法との比較

自己評価

自分の回答

気づき・メモ