始めに
FFTについて、勉強する機会があったため備忘録に記事を書こうと思いました。
問題設定
離散フーリエ変換を高速に行うアルゴリズムです。 しかし、今回は次の問題を考えることにし信号処理的な背景をすべて忘れることにします。
2つの数列$A_0,\ldots,A_{N - 1}$と$B_0,\ldots,B_{M - 1}$が与えられるとする。 この時、次の数列$C_k,(0\leq k \leq N + M - 2)$を求めよ。
\begin{align} C_k = \sum_{i + j = k} A_i B_j \end{align}
実は、次の多項式を使った表現と同じ問題になります。
入力として、2つの多項式
\begin{align} F(x) = \sum_{k = 0}^{N - 1} A_k x^k, G(x) = \sum_{k = 0}^{M - 1} B_k x^k \end{align}
が与えられるとする。この時、$(F\cdot G)(x)$としてできる多項式の係数$C_k$を求めよ。$F \cdot G$を多項式の掛け算とします。
これは、多項式の $x^k $ の係数は、$i + j = k$ となる $F(x)$ の $x^i$ の係数と $G(x)$ の $x^j $の係数に関する積和になるからです。
計算の概要
多項式の積は、ナイーブに計算を行えば$O(NM)$で計算できますが、これを$O((N + M) \log{(N + M)})$で高速化に計算を行うのがFFTです。
では、どう高速に計算を行うのでしょうか。次のステップを踏むことにより高速に計算を行うことができます。 ここで$L = N + M - 1$とし、$\omega = \exp(2\pi i / L)$とします。$i$は虚数単位です。
- 多項式から$F(\omega^0), F(\omega^1), \ldots, F(\omega^{L - 1})$と、$G(\omega^0), G(\omega^1), \ldots, G(\omega^{L - 1})$を計算する。
- $F(\omega^0), F(\omega^1), \ldots, F(\omega^{L - 1})$と、$G(\omega^0), G(\omega^1), \ldots, G(\omega^{L - 1})$から、$(F\cdot G)(\omega^0), (F\cdot G)(\omega^1), \ldots, (F\cdot G)(\omega^{L - 1})$を計算する。
- $(F\cdot G)(\omega^0), (F\cdot G)(\omega^1), \ldots, (F\cdot G)(\omega^{L - 1})$から、$F\cdot G$（の係数）を計算する。
ここで、1, 3のステップで、FFTと逆FFTを用いることができます。ここの計算量は$O(L\log L)$となります。 また、2のステップは$O(L)$で計算が可能です。
ステップ2に関しては$(F\cdot G)(\omega^k) = F(\omega^k)\cdot G(\omega^k)$と計算が可能です。これは、次の式から計算が可能です。
\begin{align} (F\cdot G)(\omega^k) &= \sum_{k = 0}^{L - 1}\sum_{i + j = k}A_i B_j \omega^k\\ &= \sum_{k = 0}^{L - 1}\sum_{j = 0}^kA_j \omega^j B_{k - j} \omega^{k - j}\\ &= \sum_{j = 0}^{L - 1} \sum_{k = j}^{L - 1} A_j \omega^j B_{k - j} \omega^{k - j}\\ &= \sum_{j = 0}^{L - 1} A_j \omega^j \sum_{k = j}^{L - 1} B_{k - j} \omega^{k - j}\\ &= \sum_{j = 0}^{L - 1} A_j \omega^j \sum_{l = 0}^{L - j - 1} B_l\omega^l\\ &= \sum_{j = 0}^{L - 1} A_j \omega^j \sum_{l = 0}^{L - 1} B_l\omega^l \end{align}
最後の行では、
\begin{align} \sum_{j = 0}^{L - 1} A_j \omega^j \sum_{l = L - j}^{L - 1} B_l\omega^l \end{align}
という値を加算しています。この加算した値について、$L = N + M - 1$と設定したため、$j \geq N$または$l \geq M $のいずれかが成り立ちます。 ここで元の多項式の定義から、$A_j$または$B_l$が$0$になります。
FFTについて
残っているステップとして、 多項式から$F(\omega^0), F(\omega^1), \ldots, F(\omega^{L - 1})$を計算する部分が残っています。 ここで計算の都合上$L = 2^h$とします。 これは必要なら$L$を大きくすることで、今までの議論の一般性を失わせないです。
FFTでは、これを分割統治法によって解くことができます。
$L=1$の時は、多項式の係数$A_0$が答えです。以下では、$L>1$の時を考えます。
次のように多項式を定義します。
\begin{align} E(x) = \sum_{j = 0}^{L/2 - 1}A_{2j}x^j\\ O(x) = \sum_{j = 0}^{L / 2 - 1} A_{2j+1}x^j\\ \end{align}
これに対して、$E(\omega^{0}), E(\omega^{2}),\ldots, E(\omega^{L - 2})$と$O(\omega^{0}), O(\omega^{2}),\ldots, O(\omega^{L - 2})$が計算できたと仮定します。
これらの2つの数列を用いて、$F(\omega^0), F(\omega^1), \ldots, F(\omega^{L - 1})$を計算することを考えます。 実は、この計算は次のように計算できます。
\begin{align} F(\omega^k) = \left\{\begin{matrix} E(\omega^{2k}) + \omega^kO(\omega^{2k}) & (k\leq L /2 - 1\text{のとき})\\ E(\omega^{2k - L}) - \omega^{k - L/2}O(\omega^{2k - L}) & (k\geq L /2\text{のとき})\\ \end{matrix}\right. \end{align}
これは下の式によって確かめることができます。
$k \leq L / 2 - 1$の時、
\begin{align} F(\omega^k) &= \sum_{j = 0}^{L - 1} A_j (\omega^k)^j\\ &= \sum_{j = 0}^{L/2 - 1}A_{2j}(\omega^k)^{2j} + \sum_{j = 0}^{L / 2 - 1} A_{2j+1}(\omega^k)^{2j + 1}\\ &= \sum_{j = 0}^{L/2 - 1}A_{2j}(\omega^{2k})^j + \omega^k \sum_{j = 0}^{L / 2 - 1} A_{2j+1}(\omega^{2k})^j\\ &= E(\omega^{2k}) + \omega^k O(\omega^{2k}) \end{align}
となります。
また、 $k \geq L / 2 $の時、
\begin{align} F(\omega^k) &= \sum_{j = 0}^{L - 1} A_j (\omega^k)^j\\ &= \sum_{j = 0}^{L/2 - 1}A_{2j}(\omega^k)^{2j} + \sum_{j = 0}^{L / 2 - 1} A_{2j+1}(\omega^k)^{2j + 1}\\ &= \sum_{j = 0}^{L/2 - 1}A_{2j}(\omega^{2k})^j + \omega^k \sum_{j = 0}^{L / 2 - 1} A_{2j+1}(\omega^{2k})^j\\ &= \sum_{j = 0}^{L/2 - 1}A_{2j}(\omega^k)^{2j}\omega^{-Lj} -\omega^{-L/2} \omega^k \sum_{j = 0}^{L / 2 - 1} A_{2j+1}(\omega^k)^{2j}\omega^{-Lj}\\ &= \sum_{j = 0}^{L/2 - 1}A_{2j}(\omega^{2k - L})^j -\omega^{k - L/2} \sum_{j = 0}^{L / 2 - 1} A_{2j+1}(\omega^{2k - L})^j\\ &= E(\omega^{2k - L}) - \omega^{k - L/2} O(\omega^{2k - L}) \end{align}
となります。ここで、$\omega^{-L/2} = -1$であり$\omega^{-L}=1$を用いました。
まとめると、$L > 1$の時は次のように計算が可能です。
逆FFT
逆FFTは、$(F\cdot G)(\omega^0), (F\cdot G)(\omega^1), \ldots, (F\cdot G)(\omega^{L - 1})$から、$F\cdot G$（の係数$C_k$）を計算することでした。
実は、この計算はFFTと同じ計算によって解くことができます。 多項式を
\begin{align} \tilde{F}(x) = \sum_{j = 0}^{L - 1} (F\cdot G)(\omega^j) x^j \end{align}
と定義します。 この多項式に対して$\tilde{\omega} := \omega^{-1}$としてFFTを計算すると、 $L(F\cdot G)(x)$（の係数）を計算が可能です。$\tilde{\omega}$を用いてFFTができることは上の議論を追うとわかります。
したがって、後は得られた答えに対して、$L$で割った値が求めたい多項式$(F\cdot G)(x)$となっています。
逆FFTがFFTの計算と同じであることは、次の計算で確かめられます。
\begin{align} \tilde{F}(\omega^{-k}) &= \sum_{j = 0}^{L - 1} (F\cdot G)(\omega^j) (\omega^{-k})^j\\ &= \sum_{j = 0}^{L - 1} \left(\sum_{l = 0}^{L - 1}C_l(\omega^j)^l\right) (\omega^{-k})^j\\ &= \sum_{j = 0}^{L - 1}\sum_{l = 0}^{L - 1}C_l\omega^{(l - k)j}\\ &= \sum_{j = 0}^{L - 1}C_k 1 ^j\\ &= LC_k \end{align}
ここで、下から2行目に関する和は、次のことより導けます。
\begin{align} \sum_{j = 0}^{L - 1} \omega^{kj} = \left\{\begin{matrix} L& (k = 0\text{の時})\\ 0 & (\text{otherwise}) \end{matrix}\right. \end{align}
Cooler-TurkeyのFFTについて
以上で積和についての問題を解くことができました。 しかし、FFTの計算において$E,O$という2つの多項式を定義する必要がありました。 これはナイーブな実装上メモリ確保が必要となり、ここをなくしたいという気持ちが残ります。 すなわち、サイズ$L(=2^h)$の配列$A$を用いて、$F$のFFTを計算したいです。
2つの多項式をまとめる際には、
\begin{align} F(\omega^k) = \left\{\begin{matrix} E(\omega^{2k}) + \omega^kO(\omega^{2k}) & (k\leq L /2 - 1\text{のとき})\\ E(\omega^{2k - L}) - \omega^{k - L/2}O(\omega^{2k - L}) & (k\geq L /2\text{のとき})\\ \end{matrix}\right. \end{align}
という更新式に従って更新を行いました。これを見ると、2つの値を使って2つの値を導出する形になっています。 具体的には、すべての$k (0\leq k\leq L/2 - 1)$に対して$F(\omega^k), F(\omega^{k + L/2})$を
\begin{align} F(\omega^k) &= E(\omega^{2k}) + \omega^k O(\omega^{2k})\\ F(\omega^{k + L / 2}) &= E(\omega^{2k}) - \omega^kO(\omega^{2k}) \end{align}
と求めることになります。 また、用いる2つの値は1度使われるとその後は用いることがないです。
したがって、うまいこと計算の順番を考えることで、$O,E$用に新たに配列を確保することなく使いまわすことができます。
ここでは、分割統治法においてマージする段階で簡単になるように配置を考えます。 具体的には、$E,O$のFFTの結果が
\begin{align} A[0:L/2 - 1] = [ E(\omega^{0}), E(\omega^{2}), \ldots, E(\omega^{L - 2})]\\ A[L/2:L - 1] = [ O(\omega^{0}), O(\omega^2), \ldots, O(\omega^{L - 2})] \end{align}
と配置されるように配列を分割することを考えます。 一旦このように配置されれば、次のように$F$のFFTを計算することが可能です。
int b = L / 2; std::complex<double> z =1; for (int j = 0; j < L / 2; j++) { auto s = A[i]; auto t = A[i + b] * z; A[k] = s + t; A[k + b] = s - t; z *= omega; }
さて、上は分割統治法において、最後にマージする部分を考えました。 次は、上のような$O,E$の分割が再帰的にできると仮定して、$A$がいくつか分割されている場合のコードを考えます。ここでは仮に$A$が$2^{k+1}$個に分割されており、$2^{k}$個にまとめることを考えます。
この処理を行うコードは、
int b = L / 2**(k+1); for (int j = 0; j < b * 2; j += 2*b) { std::complex<double> z = 1; for (int k = 0; k < b; k++) { auto s = A[j + k]; auto t = A[j + k + b] * z; A[j + k] = s + t; A[j + k + b] = s - t; z *= omega; } }
となります。上のfor文においてomegaが共通する部分の計算を、一括で行うと、
int b = L / 2**(k+1); std::complex<double> z = 1; for (int j = 0; j < b; j++) { for (int k = 0; k < b*2; k += 2*b) { auto s = A[i]; auto t = A[i + b] * z; A[k] = s + t; A[k + b] = s - t; } z *= omega; }
となります。
最終的に考えないといけない部分は、$F$の係数をどのように配列$A$に配置しなおすかということです。 ここで初めの多項式$F(x)$の係数を$A_0,A_1,\ldots, A_{L - 1}$とする。 これを分割していくと、
\begin{align} \{A_0,A_2,\ldots, A_{L - 2}\}, &\{A_1,A_3,\ldots, A_{L - 1}\}\\ \{A_0,A_4,\ldots, A_{L - 4}\}, \{A_2,A_6,\ldots, A_{L - 2}\},& \{A_1,A_5,\ldots, A_{L - 3}\}, \{A_3,A_7,\ldots, A_{L - 1}\}\\ \vdots& \end{align}
となります。 係数の添字の２進数表記を見ると、1の位が$0,1$かで分け、次は2の位が$0,1$かで分け、……という形になっています。 具体的には、$L = 8$の場合で考えると、次のように係数が分かれています。
\begin{align} \{000, 010, 100, 110 \},& \{001, 011, 101, 111\}\\ \{000, 100\}, \{010, 110 \},& \{001, 101\}, \{011, 111\} \end{align}
この配置する前のindexと配置した後のindexはちょうどbit反転の関係になっています。具体的に$L = 8$の場合では、
0 = 000 -> 000 1 = 001 -> 100 2 = 010 -> 010 3 = 011 -> 110 4 = 100 -> 001 5 = 101 -> 101 6 = 110 -> 011 7 = 111 -> 111
このように係数を配置しなおすコードを書くと
for (int i = 0; i < L; i++) { int j = 0; for (int k = 0; k < h ;k++) j |= (i >> k & 1) << (h - 1 - k); if (i < j) swap(A[i], A[j]); }
となります。
以上で、分割が終わり、先ほどのマージのコードを、分割の大きさが2の状態から計算を行えば求まります。
終わりに
本来は、もう少し図を入れたほうがわかりやすいだろうけど、今回はこれで終わりにしておきます。
参考
残りは気分が向けば書きます。
