Day 043-Q3 — 多項式合成 f(g(x)) mod x^N

2026-05-27 赤色 Master / Phase 8+ ★★★★★★★★★ FPS Composition + Baby-step Giant-step

問題

形式的冪級数 $f(x) = \sum_{k=0}^{N-1} a_k x^k$ と $g(x) = \sum_{k=0}^{N-1} b_k x^k$（$b_0 = 0$）が与えられる。 $h(x) = f(g(x)) \bmod x^N$ を求めよ。係数は $p = 998244353$ で mod を取ること。

制約

$1 \le N \le 1000$
$b_0 = 0$（合成可能条件）
$0 \le a_i, b_i < p$
時間制限: 2sec / メモリ: 256MB

入出力例

入力例 1

4
1 1 1 1
0 1 1 0

出力例 1

1 1 3 4

$f(x) = 1+x+x^2+x^3$, $g(x) = x+x^2$, $f(g(x)) \bmod x^4$ の係数列

概念図: Baby-step Giant-step 多項式合成

BSGS での多項式合成: K = √N Baby Steps (K個) g^0 = 1 g^1 = g(x) g^2 = g(x)² ... g^{K-1} 各 O(N log N) で計算: 合計 O(K·N log N) Giant Step G = g^K (一度だけ計算) G^0 = 1 G^1 = G(x) ... G^{K-1} 各 O(N log N): 合計 O(K·N log N) 組み合わせ $f(g) = \sum_{j=0}^{K-1} P_j(g) \cdot G^j$ ここで $P_j(g) = \sum_{i=0}^{K-1} a_{jK+i} \cdot g^i$（Baby stepの線形結合） 計算量 $O(N^{1.5} \log N)$ 愚直: $O(N^3)$

ヒント（段階的開示）

ヒント1: 方向性
$h(x) = f(g(x)) = \sum_{k=0}^{N-1} a_k g(x)^k$。$b_0 = 0$ なので $g^k \bmod x^N$ は $k$ 次から始まる。 愚直: $g^k$ を順次求めて $a_k g^k$ を加算 → $O(N^3)$（$N \le 1000$ なら通る）。
ヒント2: アプローチ
Baby-step Giant-step: $K = \lfloor\sqrt{N}\rfloor$ として
  • Baby: $g^0, g^1, ..., g^{K-1}$ を事前計算 ($K$ 個)
  • Giant: $G = g^K$ を計算し、$G^0, G^1, ...$ を準備 ($K$ 個)
  • $f = \sum_{j} (\sum_i a_{jK+i} g^i) G^j$ と分解して $O(K^2 \cdot N \log N)$ で計算
ヒント3: 実装骨格 
MOD = 998244353

def poly_mul_mod(p, q, N):
    res = [0] * N
    for i in range(N):
        if p[i] == 0: continue
        for j in range(N - i):
            res[i+j] = (res[i+j] + p[i] * q[j]) % MOD
    return res

# 愚直 O(N^3)
gk = [0] * N; gk[0] = 1  # g^0 = 1
result = [0] * N
for k in range(N):
    if a[k]: result = [r + a[k]*g for r, g in zip(result, gk)]...
    gk = poly_mul_mod(gk, b, N)

模範解答 (Python)

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

def solve():
    N = int(input())
    a = list(map(int, input().split()))
    b = list(map(int, input().split()))

    def poly_mul_mod(p, q):
        res = [0] * N
        for i in range(N):
            if p[i] == 0: continue
            for j in range(N - i):
                res[i+j] = (res[i+j] + p[i] * q[j]) % MOD
        return res

    # 愚直 O(N^3): N<=1000 なら PyPy で通る
    gk = [0] * N
    gk[0] = 1  # g^0 = 1
    result = [0] * N

    for k in range(N):
        if a[k] != 0:
            for i in range(N):
                result[i] = (result[i] + a[k] * gk[i]) % MOD
        if k < N - 1:
            gk = poly_mul_mod(gk, b)

    print(*result)

solve()

Step-by-Step 解説

1合成の定義と有限性
$h = f(g) = \sum_k a_k g^k$。$b_0 = 0$ なので $g^k$ の最低次数は $k$。$k \ge N$ では mod $x^N$ が 0 になるため、$k = 0$ から $N-1$ までの有限和で計算できる。
2愚直法の実装
gk = g^k を前ステップの g^{k-1} に g を掛けて順次生成。各ステップで a[k] * gk を result に加算。$O(N)$ ステップ × $O(N^2)$ 乗算 = $O(N^3)$。
3Baby-step Giant-step の概要
$K = \lceil\sqrt{N}\rceil$ として Giant step $G = g^K$ を事前計算。$f = \sum_j P_j(g) \cdot G^j$ と分解し、$P_j$ の計算に Baby step の結果を使う。各乗算を NTT $O(N \log N)$ にすれば全体 $O(N^{1.5} \log N)$。
4NTT との組み合わせ
$N > 1000$ では多項式乗算を NTT に置き換える。Python では外部ライブラリか PyPy の高速実装が必要。Library Checker の問題では $N \le 2 \times 10^4$ が標準的な制約。

計算量

愚直: $O(N^3)$（$N \le 1000$）
BSGS + 愚直乗算: $O(N^{2.5})$
BSGS + NTT: $O(N^{1.5} \log N)$
空間: $O(N)$（多項式 K 個分）

よくあるミス

ミス原因正しい書き方
b[0] != 0 で誤った合成非合成可能な入力事前に assert b[0] == 0
gk の更新を忘れるループ外で更新漏れループ末尾で gk = poly_mul_mod(gk, b)
係数の mod 取り忘れ中間値がオーバーフロー各演算後に % MOD
N=1000 で O(N^3) TLEPython の定数が大きいPyPy 使用 or BSGS に切り替え

次のステップ

  • 発展問題: $N \le 2 \times 10^4$ の合成（Brent-Kung $O(N^{1.5} \log N)$）
  • 類題: Library Checker "Composition of Formal Power Series"
  • 応用: 関数の反復合成 $f^k(x)$、EGF の変数変換

自己評価