FFTとNTTとFMTと

円周率の計算や暗号技術、競技プログラミングなどでは通常の変数型では扱えない長さの数を使った計算をすることがあり、特に掛け算をする際に筆算よりも効率的に計算する方法として FFT(高速フーリエ変換) が使われることがよく知られています。また、そのバリエーションとして NTT(数論変換) や FMT(高速剰余変換) があることも知られています。この記事ではその3つの違いについてお伝えしようと思います。

DFT と畳み込み乗算、FFT と多倍長乗算

まずは FFT を掛け算に組み込める前提となる、DFT(離散フーリエ変換)と畳み込み乗算の関係について知る必要があります。長さ $n$ の数列 $\{a_j\}$ の DFT は $\omega=\exp(-2\pi i/n)$ として

A_k = \sum_{j=0}^{n-1} a_j \omega^{jk},\quad
a_j = n^{-1} \sum_{k=0}^{n-1} A_k \omega^{-jk}

で定義されます。そして同じ長さの 2 数列 $\{a_j\}$、$\{b_j\}$ について

1. $\{a_j\}$、$\{b_j\}$ のそれぞれに DFT をかける
2. 同じインデックスの項を掛け合わせる
3. 逆 DFT をかける

という手順を踏むことで巡回畳み込み乗算

\begin{eqnarray}
c_l &=& n^{-1}\sum_{k=0}^{n-1}\left(\sum_{j=0}^{n-1}a_j\omega^{jk}\right)
        \left(\sum_{j'=0}^{n-1}b_{j'}\omega^{j'k}\right)\omega^{kl}\\
    &=& n^{-1}\sum_{k=0}^{n-1}\sum_{j=0}^{n-1}\sum_{j'=0}^{n-1}a_jb_{j'}\omega^{(j+j'-l)k}\\
    &=& \sum_{j+j'\equiv l \bmod n}a_jb_{j'}
\end{eqnarray}

を求めることができます。あとで応用されるポイントとして

\sum_{k=0}^{n-1}\omega^{jk} =
\begin{cases}
0 & (j \neq 0)\\ n & (j = 0)
\end{cases}

を覚えておいてください。巡回畳み込み乗算の式の 2 行目から 3 行目への移行で適用されています。また、$\{a_j\}$、$\{b_j\}$ の上位部分に適切な数の 0 が入っていれば巡回部分を 0 にすることができ、畳み込み乗算の結果を得られることがよく知られています。
ここで DFT と同じ結果を得られる高速化アルゴリズム FFT (後述)の計算量が $O(n\log n)$ ということから手順 1. と 3. に FFT を適用することで全体として $O(n\log n)$ の計算量で畳み込み乗算の結果を得ることができることになります。さらに畳み込み乗算の結果は正規化(繰り上がり処理を施すことで各要素の値が元の数列と同じ制約を満たすように)することで多倍長整数の乗算に応用することができます。一般に「FFT を利用して多倍長乗算を求める」と言う場合にはここまでのことを総括しています。

FFT とは

FFT についての解説は外でも多々あると思うので、ここではいくつか他との違いが見える点を掻い摘んでお伝えします。
大前提として、FFT は DFT の計算を $O(n\log n)$ の計算量で計算するアルゴリズムの名前です。$\{a_j\}$ を偶数番目と奇数番目、$\{A_k\}$ を前半と後半、という形で $n/2$ ずつの数列に仕分けることで以下のように変形します。

\begin{eqnarray}
A_k &=& \sum_{j=0}^{n/2-1} a_{2j} \omega^{2jk} + \omega^k \sum_{j=0}^{n/2-1} a_{2j+1} \omega^{2jk}\\
A_{k+n/2} &=& \sum_{j=0}^{n/2-1} a_{2j} \omega^{2jk} + \omega^{n/2}\omega^k \sum_{j=0}^{n/2-1} a_{2j+1} \omega^{2jk}\\
\end{eqnarray}

すると各 $\sum$ の部分は長さ $n/2$ の DFT になり、かつ上下に並んだ $\sum$ は同じ式なのでそれぞれ一括計算することが可能になります。また $\omega^{n/2}=-1$ となることから

\begin{eqnarray}
A_k &=& \sum_{j=0}^{n/2-1} a_{2j} \omega^{2jk} + \omega^k \sum_{j=0}^{n/2-1} a_{2j+1} \omega^{2jk}\\
A_{k+n/2} &=& \sum_{j=0}^{n/2-1} a_{2j} \omega^{2jk} - \omega^k \sum_{j=0}^{n/2-1} a_{2j+1} \omega^{2jk}\\
\end{eqnarray}

とよく見る式ができあがります。
この分割計算を再帰的に適用することで $O(n\log n)$ の計算量になることが分かります。

多倍長乗算の観点から

突然ですが、一旦 FFT から外れて多倍長乗算のお話です。一般的に多倍長整数は数列 $\{a_j\}$ を使って表現されますが、ここでは多項式の形で表現することにします。考え方の違いなだけなのでプログラムで実装する観点からの違いはありません。

\alpha(x) = \sum_{j=0}^{n-1} a_jx^j

ところで前提として基数(10進数で言うところの10)を $R$ とするとき、$0\leqq a_j < R$ が満たされるとしましょう。すると $\alpha(R)$ は表現したい多倍長整数そのものになります。そもそも多倍長乗算でやりたいことは 2 つの整数 $\alpha(R)$、$\beta(R)$ に対してその積

\gamma(R) = \alpha(R) \beta(R)

を求めることですが、多項式型の表現で見直すと一般式として

\gamma(x) = \alpha(x) \beta(x)

を求めることができれば大丈夫なはずです。そしてカラツバ法 や Toom-Cook法(リンク先は英語) はその具体的な方法ということになります。

Toom-Cook 法をざっとおさらいしてみましょう。最もよく見るパターンの Toom-Cook-3 法では $\alpha(x)$ や $\beta(x)$ が 2 次式であることを前提に、 $x=0, \pm1, \pm2$ といった 5 つの具体的な数値を代入し、得られた 5 つの積が $\gamma(0)$、$\gamma(\pm1)$、$\gamma(\pm2)$ と等しくなることを利用して 4 次多項式 $\gamma(x)$ の係数を求めていきます。
同様に多項式の次数を増やすことで $\alpha(x)$ と$\beta(x)$ の掛け算にかかる計算量は減少します (次数 $d$ に対して $O(n^{d/(d+1)})$) が、各値を代入した式を求める部分の計算量が爆発します。具体的には $2(d+1)^2$ 個の値を求めることになるので $O(d^2(n/d))$ くらい。あとは計算量の算出が面倒ですが、代入する値の最大値が $d$ なので $d-1$ 次多項式で $d^{d-1}$ となり、係数との掛け算の部分に $O(d\cdot n/d)=O(n)$ のコストがかかります。これらを全部合計すると $O(n^{d/(d+1)} + nd + nd^2)$ ということであまり分割しすぎるとかえってコスト高になります。

さて、この増えていく方の計算量を減らす方法は無いでしょうか？具体的には最後の「代入する値の最大値が $d$ なので…」辺り。ここで先程の多項式に仮代入する値は整数、というより実数に限らないので試しに $\omega^k$ という値を代入してみましょう。$k$ は数値バリエーションを得るためのパラメータだと思って下さい。

\alpha(\omega^k) = \sum_{j=0}^{d-1} a_j \omega^{jk}

あからさまな誘導ですが、結果紛れもなく DFT の形を得ることができました。つまり FFT が利用できます。しかも相異なる $k$ を $n$ 個代入すれば一次独立な式を $n$ 個得られ、逆 FFT をかけることで $\gamma(x)$ の係数が求めることができます。係数との掛け算部分のコストを減らす(絶対値を小さくする)のが目的でしたが、ついでに複数の値を纏めて求めることができる点でも計算コストが下がっていることが分かります。
ということで無理やりながら FFT に戻ってきました。多倍長の観点からといいつつ結局は元々 FFT を知ってないとこの応用を思いつくのは不可能だと思います。

NTT

と、無事 FFT に話が戻ってきたところで改めて FFT と畳み込み乗算の関係式、ならびに多倍長の観点から…の流れを見てみると、計算上必要な条件は $\omega^n=1$ となる $\omega$ (原始$n$乗根) の存在だということがわかります。その条件さえ満たせば計算空間が複素数である必要もありません。そこで登場するのが数論変換こと NTT です。知らない人には名前から伝わりにくいのですが、数論(Number Theorem; NT)とはざっくりいえば整数、もっとぶった切って言えば剰余環を対象に扱う数学のジャンルです。あ、石を投げないで……

閑話休題、NTT は $\omega^n \equiv 1 \pmod m$ を満たす 2 整数 $\omega$ と $m$ を使って $\pmod m$ の剰余環で DFT と同様の計算をすることを指します。つまり厳密には FFT ではなく DFT と対応する位置づけになります。実用的には $n$ が 2 の冪乗になるような値を選択することで DFT における FFT と同様の高速化を実現することが前提になるので FFT と並べて扱われることが殆どです。

あらためて定義ははっきりしたものの、このような数値ペア ($\omega$ と $m$) は簡単に見つけることができるのでしょうか。$m=an+1$ の形を仮定して $\omega^{n/2}\not\equiv 1$、$\omega^n\equiv 1$ を満たす $(a, \omega)$ を総当り的に探索する以下のようなコードを書いてみると結構あっさり見つかります。 Qiita っぽくするために書いたコードです。

def find_root(k, a):
    n = 2 ** k
    m = a * n + 1
    for w in range(2, 1000):
        if pow(w, n, m) == 1 and pow(w, n // 2, m) != 1:
            return w
    return None


for k in range(50, 64):
    for a in range(1, 1000):
        # 2^64 に入らない m は考えない                                                                                                                                                                             
        if a * 2 ** k + 1 >= 2 ** 64:
            break
        w = find_root(k, a)
        if w is not None:
            print("n=2^{}, m={}n+1, w={}".format(k, a, w))
            break

こんな何も考えない Python コードでも

$ time python3 find_ntt.py
n=2^50, m=7n+1, w=3
n=2^51, m=14n+1, w=19
n=2^52, m=7n+1, w=12
n=2^53, m=20n+1, w=14
n=2^54, m=10n+1, w=5
n=2^55, m=5n+1, w=3
n=2^56, m=27n+1, w=11
n=2^57, m=29n+1, w=21
n=2^58, m=54n+1, w=8
n=2^59, m=27n+1, w=87
python3 find_ntt.py  1.40s user 0.01s system 99% cpu 1.407 total

秒で見つけてくれます。なので例えば $\omega$ を固定するなり、$m$ の大きさ (計算可能な範囲に関わる) の範囲を絞るなり条件を強めて探しましょう。ちなみに $m$ が素数である($\pmod m$ が体である)必要はありません。

Schönhage-Strassen Algorithm

そういう探索をするのすら面倒なあなたには Schönhage-Strassen アルゴリズム (SSA) をおすすめします。タイトルに出てきてない変なのが出てきましたが NTT の一種です。詳細は別記事に書いていますが要点としては NTT の上記の条件のうち $\omega$ を 2 に固定し $m=2^{n/2}+1$ と定める方法です。事前計算が必要なくなった代わりに $m$ がネイティブに使われる 64bit 整数よりも大きくなるので内部演算として多倍長演算が必須になりますが、理論上計算桁数の上限を無くしてくれます。

ちなみに SSA の定義から外れますが、この方針をもっと押し進めると $\omega$ の値が何であれ $m=\omega^{n/2}+1$ の設定だけで条件が満たされることが分かります¹。 なので例えば $\omega=R$ (基数) と設定することで変換処理の主要部分を数値移動と加減算だけにすることができますし、特に $R$ に 10 の冪乗数を設定することで 10 進数のまま計算することも可能になります。

FMT

FFT と NTT(における FFT 的アルゴリズム)、似てるようでちょっと似ていないアルゴリズムですが、これらを区別せず統一的に扱うことはできないでしょうか？できます。そう、"FMT" ならね。ということで高速剰余変換 (Fast Modulo Transformation; FMT) の紹介です。FMT では多倍長整数のところで紹介したような数式を元にします。

\alpha(x) = \sum_{j=0}^{n-1}a_jx^j

これをおもむろに $(x-\omega^k)$ で割ります。ここで $\omega$ や $k$ の値が具体的に何なのかはひとまず関係ありません。記号的な演算です。

\alpha(x) = (x-\omega^k)Q(x) + \alpha(\omega^k)

こうして残った余りは FFT や NTT の変換をしたものと同様に $\alpha(\omega^k)$ になるのでこれを数式処理上の計算対象に考えることで元が複素数だったか剰余環だったかなどを気にせず理論立てられます。ということで FMT の「剰余変換」とは「多項式 $\alpha(x)$ を 1 次式 $(x-\omega)$ で割った余りに変換する」ことを意味します。NTT 的な意味での剰余は関係ありません(←これがこの記事の一番のポイント)。$\omega$ に $\exp(-2\pi i/n)$ を代入すれば DFT と同じになりますし、$\pmod m$ で $\omega$ を代入すれば NTT と同じになります。
FMT については本来 FFT や NTT の単なる拡張という観点では見つけにくい性質への応用もありますが、ここでは FFT や NTT を統一的に扱う手法以上の深入りをしないことにします。

1. $n^{-1}$ の計算が簡単なのかは別の話 ↩