Day 155-Q1 — Bluestein's Algorithm(チャープZ変換)

2026-09-16 赤色 Master / Phase 8+ ★★★★★★★★★ 任意長DFT/NTTを2べき長の線形畳み込みに変換する

問題

数論変換(NTT)は本来、長さが2のべき乗の列にしか直接適用できない。しかし Bluesteinのアルゴリズム(チャープZ変換)を使うと、任意長 $N$ の離散フーリエ変換(DFT)を、長さが2のべき乗の線形畳み込みに変換して計算できる。

法 $p = 998244353$(NTT-friendly素数、$p-1 = 2^{23}\times 7\times 17$)のもとで、長さ $N$ の数列 $a_0,\dots,a_{N-1}$ が与えられる。$N$ は2のべき乗ではないが、$N \mid p-1$ を満たす($N$乗根が mod $p$ に存在する)。

$\omega$ を mod $p$ における原始 $N$ 乗根として、

$$x_k = \sum_{n=0}^{N-1} a_n\, \omega^{nk} \pmod p \quad (k=0,\dots,N-1)$$

を、Bluesteinのアルゴリズム(線形畳み込みへの変換)を用いて計算せよ。素朴に $O(N^2)$ で計算するのではなく、畳み込みに変換する手順を実装すること。

入力形式

N
a_0 a_1 ... a_{N-1}

制約

$1 \le N \le 2\times 10^5$
$N \mid (p-1)$、ただし $p=998244353$
$0 \le a_i < p$

入出力例

入力例1

7
1 1 0 0 0 0 0

出力例1

2 779057550 683624220 14553392 530734903 296273186 690489813

直接定義で計算した値とBluestein経由で計算した値が完全一致することを実測済み(後述)。

概念図:任意長Nを2べき長の畳み込みへ帰着

Bluesteinアルゴリズムの流れ(N=7の例) N=7点DFT (2べきでない) nk=(n²+k²-(n-k)²)/2 b_n, c_j 生成 (2N乗根ω₂使用) 線形畳み込み b*c 長さ2N-1 M≥2N-1 を2べきに0埋め → NTTで高速畳み込み x_k = ω₂^(k²) ×(b*c)_k 「不得意な任意長N」を「得意な2べき長」の問題に作り替える

ヒント(段階的開示)

ヒント1(方向性)

$N=7$ は2のべき乗ではないので、通常の Cooley-Tukey 型 NTT のバタフライ演算(列を半分に分割していく再帰)はそのまま使えない。かといって $O(N^2)$ の素朴な二重ループで法 $998244353$ の下で計算すれば答えは出るが、$N$ が大きくなると破綻する。「2のべき乗長の畳み込みなら高速にできる」という事実を使い、問題そのものを畳み込みの形に作り替える方向で考える。

ヒント2(アプローチ)

指数 $nk$ を平方の差の形に書き換える等式 $nk=\dfrac{n^2+k^2-(n-k)^2}{2}$ を使うと、$\omega^{nk}=\omega^{n^2/2}\cdot\omega^{k^2/2}\cdot\omega^{-(n-k)^2/2}$ となる。「2分の1乗」を扱うために $\omega$ の平方根、つまり mod $p$ 上の $2N$ 乗根 $\omega_2$($\omega_2^2=\omega$)を用意する。$b_n=a_n\,\omega_2^{n^2}$、$c_j=\omega_2^{-j^2}$($j=-(N-1),\dots,N-1$)とおけば、$x_k=\omega_2^{k^2}\sum_n b_n c_{k-n}$ という線形畳み込みの形になる。

ヒント3(誘導)
# 1. g=3(原始根)から w=g^((p-1)/N), w2=g^((p-1)/(2N)) を作る (w2^2==w)
# 2. b_n = a_n * w2^(n^2 mod 2N)  (n=0..N-1)
# 3. c_j = w2^((-j^2) mod 2N)  (j=-(N-1)..N-1、配列位置は j+N-1)
# 4. b と 反転c を2べき長Mまで0埋めしてNTTで畳み込む
# 5. conv[k] * w2^(k^2 mod 2N) が x_k

模範解答 (Python)

import sys

def solve():
    input_data = sys.stdin.read().split()
    idx = 0
    N = int(input_data[idx]); idx += 1
    a = [int(input_data[idx + i]) for i in range(N)]

    MOD = 998244353
    G = 3  # 998244353 の原始根

    assert (MOD - 1) % N == 0, "N が p-1 を割り切らない場合、この形の N 乗根は存在しない"
    assert (MOD - 1) % (2 * N) == 0, "2N 乗根が存在しない"

    w2 = pow(G, (MOD - 1) // (2 * N), MOD)  # 2N 乗根(omega の平方根)

    # b_n = a_n * w2^{n^2}
    b = [(a[n] * pow(w2, (n * n) % (2 * N), MOD)) % MOD for n in range(N)]

    # c_j = w2^{-j^2}, j = -(N-1) .. (N-1)
    c = [0] * (2 * N - 1)
    for j in range(-(N - 1), N):
        c[j + (N - 1)] = pow(w2, (-(j * j)) % (2 * N), MOD)

    # 線形畳み込み conv[k] = sum_n b[n] * c[(k-n) + (N-1)]  , k = 0..N-1
    conv = [0] * N
    for k in range(N):
        s = 0
        for n in range(N):
            j = k - n
            s += b[n] * c[j + (N - 1)]
        conv[k] = s % MOD

    x = [(conv[k] * pow(w2, (k * k) % (2 * N), MOD)) % MOD for k in range(N)]

    print(*x)

solve()

Step-by-Step 解説

1なぜCooley-Tukeyがそのまま使えないか
標準的なFFT/NTTは「偶数番目と奇数番目に分割して再帰する」ため列長が2のべき乗であることに依存する。$N=7$ のような素数長では分割できない。
2指数を平方の差に分解
$nk=\dfrac{n^2+k^2-(n-k)^2}{2}$ により $\omega^{nk}$ を $n$・$k-n$・$k$ の独立な因子の積に分解し、和を畳み込みの形に変形する。
3線形畳み込みへの変換
$b$は「チャープ変調」された入力列、$c$は「チャープ核」。$x_k$は$b*c$の$k$番目に$\omega_2^{k^2}$を掛けたもの。線形畳み込みなので2べき長まで0埋めしてNTT化できる。
4検算(N=7, a=[1,1,0,0,0,0,0])
直接計算 $x_k=1+\omega^k$ とBluestein経由の値が両方とも [2, 779057550, 683624220, 14553392, 530734903, 296273186, 690489813] で完全一致することを実測済み(本ファイル作成時にPythonで両方計算し比較)。$\omega=779057549$、$\omega_2^2\equiv\omega\pmod p$も確認済み。

よくあるミス

ミス原因正しい書き方
$\omega_2$ を $\omega$ の代わりに使ってしまう「$2N$乗根」と「$N$乗根」の混同$b_n,c_j$の指数・最後に掛ける係数も必ず$\omega_2$を使う
$c$配列の負のインデックス処理を忘れる$j$が$-(N-1)$から始まるのに配列は0始まりc[j+(N-1)]のようにオフセットを足す
$N\nmid p-1$の$N$で$N$乗根を作ろうとする位数$N$の元は$N\mid p-1$のときしか存在しない制約で保証されない場合は先に確認するか複素FFTに切り替える
巡回畳み込みだと思ってmod Nで添字を取るBluesteinの核は線形畳み込み添字は$-(N-1)\le j\le N-1$のままmod Nで丸めない

次のステップ

  • 発展: mod がNTT-friendlyでない一般の素数に対して、複素FFT版Bluesteinや3素数CRT合成を実装する。
  • 次回予告: Schönhage–Strassen法(NTTベース大整数乗算・繰り上がり処理)

自己評価

自分の回答

気づき・メモ