競プロ用ライブラリを作る Advent Calendar 2024の8日目です.
FFTって?
Fast Fourier Transformの略で, 離散フーリエ変換というものを高速に行えます.
離散フーリエ変換とは, $N$個の数値$a_0,a_1,...a_{N-1}$が与えられたとき
\hat a_t=\sum_{k=0}^{N-1}a_k\exp\left(\frac{2\pi i}N kt\right)
としたときの$\hat a_0, \hat a_1, ..., \hat a_{N-1}$を計算することです. $i$は虚数単位です.
愚直にやると$O(N^2)$かかりますが, Nが2の冪のときにこれが$O(N\log N)$まで落ちます.
仕組み
$N$が偶数だと仮定し, 上の式を変形します.
\begin{eqnarray}
\hat a_t&=&\sum_{k=0}^{N-1}a_k\exp\left(\frac{2\pi i}Nkt\right) \\
&=&\sum_{k=0}^{N-1}a_k\exp\left(\frac{2\pi i}Nt\right)^k \\
&=&\left(\sum_{k=0}^{\frac N2-1}a_{2k}\exp\left(\frac{2\pi i}Nt\right)^{2k}\right)+\left(\sum_{k=0}^{\frac N2-1}a_{2k+1}\exp\left(\frac{2\pi i}N t\right)^{2k+1}\right) \\
&=&\left(\sum_{k=0}^{\frac N2-1}a_{2k}\exp\left(\frac{2\pi i}Nt\right)^{2k}\right) + \exp\left(\frac{2\pi i}Nt\right)\left(\sum_{k=0}^{\frac N2-1}a_{2k+1}\exp\left(\frac{2\pi i}Nt\right)^{2k}\right) \\
\end{eqnarray}
やったことは偶数番目の項と奇数番目の項に分けた感じです.
さて$b_k = a_{2k}, c_k = a_{2k+1}, M = N\div 2$と置いてみます. $N$は偶数と仮定を置いたので$M$は必ず整数ですね.
\begin{eqnarray}
\hat a_t&=&\left(\sum_{k=0}^{M-1}b_k\exp\left(\frac{2\pi i}Nt\right)^{2k}\right)+\exp\left(\frac{2\pi i}Nt\right)\left(\sum_{k=0}^{M-1}c_k\exp\left(\frac{2\pi i}Nt\right)^{2k}\right) \\
&=&\left(\sum_{k=0}^{M-1}b_k\exp\left(\frac{2\pi i}Mkt\right)\right)+\exp\left(\frac{2\pi i}Nt\right)\left(\sum_{k=0}^{M-1}c_k\exp\left(\frac{2\pi i}Mkt\right)\right) \\
\end{eqnarray}
ここで, 2つの$\sum$記号に注目すると, 最初の離散フーリエ変換の式の$x, N$が$b, c, M$に置き換わったものだと分かります.
高速フーリエ変換では, これを「半分のサイズの離散フーリエ変換2つ」にしてしまいます.
とりあえず$b, c$を離散フーリエ変換したものを$\hat b, \hat c$と書くことにすると,
\begin{eqnarray}
\hat a_t&=&\hat b_t+\exp\left(\frac{2\pi i}Nt\right)\hat c_t \\
\end{eqnarray}
とすっきりした形になります.
この式の$t$ですが, 求めたいのは$\hat a_0, \hat a_1, ..., \hat a_{N-1}$なので$0\le t\le N$の範囲を動きます.
しかし, 半分のサイズの離散フーリエ変換では$N$が$M$に置き換わっていることに気を付けると, $b, c$をそれぞれで離散フーリエ変換しても求まるのは$\hat b_0, \hat b_1, ..., \hat b_{\color{red}M-1}$と$\hat c_0, \hat c_1, ..., \hat c_{\color{red}M-1}$で, $M\le t\lt N$の範囲の値が欠けてしまっていることが分かります.
どうするかというと, 実は$\hat b_t = \hat b_{t+M}, \hat c_t = \hat c_{t+M}$が成り立つので$0\le t\lt M$の範囲がそのまま$M\le t\lt 2M=N$の範囲の値として使えます.
というのも, $k$が整数なら$\exp(x) = \exp(x+2\pi i k)$が成り立つので
\begin{eqnarray}
\hat b_t &=& \sum_{k=0}^{M-1}b_k\exp\left(\frac{2\pi i}Mkt\right) \\
&=& \sum_{k=0}^{M-1}b_k\exp\left(\frac{2\pi i}Mkt+2\pi ik\right) \\
&=& \sum_{k=0}^{M-1}b_k\exp\left(\frac{2\pi i}Mk(t + M)\right) \\
&=& \hat b_{t+M}
\end{eqnarray}
という感じになるからです.
こんな感じで, 離散フーリエ変換を半分のサイズの離散フーリエ変換$2$つにわけることが出来ました.
その「半分のサイズの離散フーリエ変換」に, またこの操作をすると$1/4$のサイズの離散フーリエ変換$4$つになります.
これを繰り返して, ある程度まで問題が小さくなったら愚直に離散フーリエ変換をしてやって, 元に戻すとで答えが復元できます.
これで計算量は$O(N\log N)$です.
この方法では半分の問題にするのには$N$が偶数である必要があるので, $N$を2の冪などに揃えてやるのが一般的です.
畳み込み
競プロではだいたい畳み込みの計算に使われるので, その解説もします.
畳み込みとは, ここでは多項式の掛け算のことって考えて問題ないです.
さて, 最初の
\hat a_t=\sum_{k=0}^{N-1}a_k\exp\left(\frac{2\pi i}N kt\right)
の式ですが, $f(x) = \sum_{k=0}^{N-1}a_kx^k$ という多項式を考えると$\hat a_t=f\left(\exp\left(\frac{2\pi i}Nt\right)\right)$になってますね.
同じように数列$b_0, b_1, ..., b_{N-1}$の離散フーリエ変換$\hat b$は$g(x) = \sum_{k=0}^{N-1}b_kx^k$という多項式を考えると$\hat b_t=g\left(\exp\left(\frac{2\pi i}Nt\right)\right)$になります.
さて, 更に数列$c_t = \hat a_t\hat b_t$を考えます. $c_t = f\left(\exp\left(\frac{2\pi i}Nt\right)\right)g\left(\exp\left(\frac{2\pi i}Nt\right)\right)$ということですが, 多項式$h(x) = f(x)g(x)$を考えると$c_t = h\left(\exp\left(\frac{2\pi i}Nt\right)\right)$になります.
$h$の係数を離散フーリエ変換したものがちょうど$c$になっていることが分かります.
離散フーリエ変換されたものを元の多項式に戻せれば, 多項式$f$と$g$の積$h$が分かりますね. 実際フーリエ逆変換というのがあって
a_t=\frac1N\sum_{k=0}^{N-1}\hat a_k\exp\left(-\frac{2\pi i}N kt\right)
となっています. さっきのフーリエ変換とほぼ同じ形なので離散フーリエ変換と同じような議論がで$O(N\log N)$で計算できます.
というわけで2つの多項式を$O(N\log N)$で離散フーリエ変換し, $O(N)$で積を求め, $O(N\log N)$でフーリエ逆変換するので, 合わせて$O(N\log N)$で多項式の積を求めることが出来ました.
実装するときの注意として, $h$に$N$次以上の項があると溢れて変な結果になるので, $N$を$f$の次数+$g$の次数より大きくするようにしましょう.
実装
ひたすら計算したので, それをソースコードに落としてやればいいです.
実装は複素数ではなく剰余類環上で計算します. 沢山でてきた$\exp\left(\frac{2\pi i}N\right)$みたいな式も「$N$乗したら$1$になるけど$2$乗, $3$乗, ..., $N-1$乗では$1$にならない」みたいな性質が重要で, そういう性質の値を持つ剰余類環でも同じことが計算できてしまう訳です.
まとめ
いかがでしたか?
明日は創作データ構造します.