Day 011-Q2 | 形式的冪級数 (FPS) 基礎

問題

数列 $A = [a_0, a_1, \ldots, a_{n-1}]$ と $B = [b_0, b_1, \ldots, b_{n-1}]$ が与えられる。これらを多項式 $f(x) = \sum_{i=0}^{n-1} a_i x^i$、$g(x) = \sum_{i=0}^{n-1} b_i x^i$ とみなす。

$h(x) = f(x) \cdot g(x) \pmod{x^n}$ の係数列 $[h_0, h_1, \ldots, h_{n-1}]$ を 素数 $p = 998244353$ を法として求めよ。

入力形式

n
a_0 a_1 ... a_{n-1}
b_0 b_1 ... b_{n-1}

制約

$1 \le n \le 2 \times 10^5$

$0 \le a_i, b_i < p$

入出力例

入力例 1

4
1 2 3 4
5 6 7 8

出力例 1

5 16 34 60

説明: $(1+2x+3x^2+4x^3)(5+6x+7x^2+8x^3) \pmod{x^4}$

$x^0$: $1\cdot5 = 5$
$x^1$: $1\cdot6+2\cdot5 = 16$
$x^2$: $1\cdot7+2\cdot6+3\cdot5 = 34$
$x^3$: $1\cdot8+2\cdot7+3\cdot6+4\cdot5 = 60$

ヒント (段階的開示)

ヒント1: 方向性

$O(n^2)$ の素朴な畳み込みは $n=2\times10^5$ では TLE。NTT（数論変換）を使うと $O(n \log n)$ で解ける。

ヒント2: アプローチ

$p = 998244353 = 119 \times 2^{23} + 1$ は NTT に都合のよい素数（NTT 素数）。原始根 $g = 3$ を使い FFT を整数演算で実行する。

ヒント3: 誘導

MOD = 998244353
# 原始根 g = 3, g^{(p-1)/2^k} が k-th 回転因子
# NTT の実装:
# 1. 長さを2の冪に拡張
# 2. ビット反転置換
# 3. バタフライ演算（乗算はmodで）

模範解答 (Python)

import sys
input = sys.stdin.readline

MOD = 998244353

def power(a, b, mod):
    res = 1
    a %= mod
    while b > 0:
        if b & 1:
            res = res * a % mod
        a = a * a % mod
        b >>= 1
    return res

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 = power(3, (MOD - 1) // length, MOD)
        if invert:
            w = power(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 = power(n, MOD - 2, MOD)
        for i in range(n):
            a[i] = a[i] * inv_n % MOD

def multiply(f, g, size):
    result_size = size
    n = 1
    while n < len(f) + len(g):
        n <<= 1
    fa = f[:] + [0] * (n - len(f))
    ga = g[:] + [0] * (n - len(g))
    ntt(fa, False)
    ntt(ga, False)
    for i in range(n):
        fa[i] = fa[i] * ga[i] % MOD
    ntt(fa, True)
    return fa[:result_size]

def main():
    n = int(input())
    a = list(map(int, input().split()))
    b = list(map(int, input().split()))
    c = multiply(a, b, n)
    print(*c)

main()

Step-by-Step 解説

1なぜ NTT か
通常の FFT は浮動小数点誤差がある。MOD 演算が必要なときは 整数上の FFT = NTT を使う。$p = 998244353$ は $2^{23}$ が因数にあるため、長さ最大 $2^{23}$ の NTT が可能。

2ビット反転置換
再帰的バタフライを反復版にするため、インデックスをビット反転した順に並べる。

3バタフライ演算
長さ length のブロックを半分に分け「回転因子 $\omega$」をかけながら和・差を計算:

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

4逆 NTT で積を復元
正変換後に点ごとに積を取り、逆 NTT で元の空間に戻す。逆 NTT は「$\omega$ を逆元にして全体を $n$ で割る」だけ。

よくあるミス

ミス	原因	正しい書き方
MOD忘れ	バタフライ演算で爆発	毎回 `% MOD`
配列長を非2冪にする	NTT が壊れる	`while n < len(f)+len(g): n<<=1`
逆変換で `inv_n` 忘れ	係数が $n$ 倍になる	`power(n, MOD-2, MOD)` で割る

Day 011-Q2 — 形式的冪級数（FPS）基礎

問題