競プロのための高速フーリエ変換

Pocket

フーリエ変換って知ってる?」

「フーリエ変換ですか? 信号処理でもするんですか、先輩」

「いや、競プロで使えるようになりたいなって」

「ああ、畳み込みの高速化ですか。あれは高速フーリエ変換をそのまま使うだけですからね」

「勉強しようと思って色々調べたんだけど、なかなか理解できなくて」

「私もうまく説明できるかわかりませんが、お話しましょうか」

離散フーリエ変換(DFT)

フーリエ変換の仲間たち

「フーリエ変換の仲間には色々種類があるんですけど、競プロで使うのは離散フーリエ変換(DFT: Discrete Fourier Transform)ですね」

「離散ってことは、連続もあるんだよね」

「はい。ざっくり 4 種類あります。フーリエ変換では元の信号を別の信号に変換するんですが、元の信号が周期信号だと変換後は離散信号(数列)になります

「周期信号っていうのは、決まったパターンの繰り返しってことだよね」

「はい。同様に、元の信号が周期信号でない場合は、変換後は連続信号になります」

「うん。残りの 2 種類は?」

「残りというか、この 2 パターンそれぞれについて、元の信号が連続信号である場合と、離散信号である場合があります。関数か数列かってことですね」

「なるほど」

「つまり、フーリエ変換の仲間にはこの 4 種類があります」

  • 非周期的・連続関数→連続関数:フーリエ変換
  • 周期的 ・連続関数→数列  :フーリエ級数展開
  • 非周期的・数列  →連続関数:離散時間フーリエ変換
  • 周期的 ・数列  →数列  :離散フーリエ変換(DFT)

「へー。確かに、プログラムでは連続関数は扱えないから、入力も出力も離散である離散フーリエ変換しか使えないね」

「はい。フーリエ級数展開が発端なんですけど、そこから色々派生したイメージですね」

「あれ、でも数列の畳み込みに使うときって、周期的とは限らないよね?」

「離散フーリエ変換では、有限数列を繰り返した無限長の周期数列を扱います。{1, 2, 3, 4} という数列を離散フーリエ変換するときは、{…, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, …} というような無限に長い繰り返し数列をイメージします」

「ははぁ、無理やり周期数列にするのか」

「はい。とは言っても、計算には元の数列しか出てきませんよ」

離散フーリエ変換の定義

「離散フーリエ変換の定義はどうなるの?」

「定義はこちらです。$N$ は元の数列 $f$ の長さです。この式では変換後の数列 $F$ の第 $t$ 項を計算しています」

$$F[t]=\sum_{x=0}^{N-1}f[x]e^{-i\frac{2\pi tx}{N}}$$

「うーん、この意味が分からないんだよなぁ」

「意味はまた今度お話するとして、とりあえず受け入れてみましょう」

「$e^{-i\frac{2\pi}{N}}$ が 1 の原始 $N$ 乗根、つまり $N$ 乗して初めて 1 になる複素数で、その $tx$ 乗を数列 $f[x]$ に掛けて $F[t]$ を計算しているね」

「はい。書くのが面倒なので、1 の原始 $N$ 乗根の一つを $W_N=e^{-i\frac{2\pi}{N}}$ と書きましょう。こうすると、離散フーリエ変換はこうなります」

$$F[t]=\sum_{x=0}^{N-1}f[x]W_N^{tx}$$

「$F[t]$ の長さ、つまり $t$ の範囲はどうなるの?」

「$t$ は $0$ から $N-1$ までの $N$ 個です。なので、そのまま計算しようとすると、$N^2$ 回の計算が必要ですね」

「そのまま計算するのは意味が無さそうだね」

「信号処理的には意味が無くはないんですけど、競プロ的には $O(N^2)$ で DFT するのは無意味ですねー」

逆離散フーリエ変換(IDFT)

「離散フーリエ変換の結果 $F[t]$ から、元の $f[x]$ を得ることもできます。これを逆離散フーリエ変換(IDFT: Inverse DFT)と呼びます」

「へぇ、どうやるの?」

「こうします」

$$f[x]=\frac{1}{N}\sum_{t=0}^{N-1}F[t]W_N^{-tx}$$

「DFT とほとんど同じだね。掛ける係数が $W_N^{-tx}$ で DFT の逆数になっているのと、結果を $N$ で割っているところだけだね」

「はい。なので、順変換(DFT)のコードをほぼそのまま使えます。便利ですね」

「これが正しいことを証明しておこうか。DFT の定義式を使って、IDFT 結果を $f'[x]$ とすると、こうだね」

$$f'[x]=\frac{1}{N}\sum_{t=0}^{N-1}\left (\sum_{y=0}^{N-1}f[y]W_N^{ty}\right )W_N^{-tx}$$

$$=\frac{1}{N}\sum_{y=0}^{N-1}f[y]\sum_{t=0}^{N-1}W_N^{t(y-x)}$$

「ここで、$W_N$ に関する非常に重要な性質、直交性が役に立ちます」

「直交性?」

「次の性質です。この性質があるのが、DFT に $W_N$ が用いられる理由です」

$$\sum_{t=0}^{N-1}W_N^{ty}W_N^{-tx}=\sum_{t=0}^{N-1}W_N^{t(y-x)}=\begin{cases}N&x\equiv y \ \pmod{N}\\ 0 & {\rm otherwise.}\end{cases}$$

「$x\equiv y\ \pmod{N}$ の場合は、整数 $k$ を用いて $y-x=kN$ となるから、$W_N^N=1$ から、総和の中身が $W_N^{tkN}=1^{tk}=1$ になるね。これを $N$ 個足すから、$N$ だね。そうでない場合は?」

「$x\not\equiv y\ {\rm mod}\ N$ の場合は、等比級数の和の公式で計算できます。初項 $1$、公比 $W_N^{y-x}$、項数 $N$ で、公比が 1 でないので、和はこうなります」

$$\sum_{t=0}^{N-1}W_N^{t(y-x)}=\frac{1-W_N^{N(y-x)}}{1-W_N^{y-x}}=\frac{1-1}{1-W_N^{y-x}}=0$$

「なるほど。なんでこれが “直交性” なの?」

「実数ベクトルの内積は、各次元成分の積の和になりますが、複素数ベクトルの内積は、片方を複素共役にして計算されるんですね」

$$a\cdot b=\sum_{t=0}^{N-1}a_t\bar{b_t}$$

「なるほど、そうすると、今計算した総和は、2 つのベクトル $(W_N^{0y},W_N^{1y},W_N^{2y},\dots)$ と $(W_N^{0x},W_N^{1x},W_N^{2x},\dots)$ の内積ということになるのか。$W_N^k$ の複素共役は $W_N^{-k}$ だから。で、$x, y$ が異なれば内積が 0 になるから、直交していると」

「はい。元の式に戻ると、$t$ に関する総和の中身は、$y$ が $x$ に一致するときのみ $N$、それ以外では $0$ になるので、こうなります」

$$f'[x]=\frac{1}{N}\sum_{y=0}^{N-1}f[y]\sum_{t=0}^{N-1}W_N^{t(y-x)}=\frac{1}{N}Nf[x]=f[x]$$

「これで、IDFT の結果が元の数列に一致したね」

「はい。ちゃんと逆変換になっていますね」

DFT による畳み込み計算

「畳み込みというのは、次の計算だね」

$$c[k]=\sum_{j=0}^k a[j]b[k-j]$$

「そうですね。数列 $a,b$ があるときに、$c[k]$ の値を数列 $a, b$ の、添字の和が $k$ になるようなペアの積の合計とするのが畳み込みですね」

「離散フーリエ変換を使うと、この畳み込みが $O(N\log N)$ で計算できるんだよね」

「はい。正確には、DFT と IDFT が $O(N\log N)$ で計算できて、畳み込みは $O(N)$ で計算できます」

「なるほど。畳み込みはどう計算するの?」

「数列 $a,b$ の長さをそれぞれ $N_a,N_b$ とします。このとき、数列 $c$ の長さを $N$ とすると、$N=N_a+N_b-1$ となります。なので、数列 $a,b$ も末尾に $0$ をつけて、長さが $N$ になるようにしておきます」

$$a[N_a]=a[N_a+1]=\cdots=a[N]=0$$

$$b[N_b]=b[N_b+1]=\cdots=b[N]=0$$

「そして、まず畳み込み結果の $c$ を DFT してみましょう」

$$C[t]=\sum_{k=0}^{N-1}c[k]W_N^{tk}$$

$$=\sum_{k=0}^{N-1}\left (\sum_{j=0}^k a[j]b[k-j]\right )W_N^{tk}$$

$$=\sum_{k=0}^{N-1}\sum_{j=0}^k a[j]W_N^{tj} b[k-j] W_N^{t(k-j)}$$

「ここまではただ代入しただけだけど、ここからどうするの?」

「まず、総和を入れ替えます。この総和は、$0\le j,k\le N-1$ かつ $j\le k$ の範囲を全て足すことになるので、 $j$ を先に考えるとこうなります」

$$=\sum_{j=0}^{N-1}\sum_{k=j}^{N-1} a[j]W_N^{tj} b[k-j] W_N^{t(k-j)}$$

「なるほど、こうすると、$k$ に無関係な $a[j]W_N^{tj}$ を内側のシグマの外に持ってこられるね」

$$=\sum_{j=0}^{N-1}a[j]W_N^{tj}\sum_{k=j}^{N-1} b[k-j] W_N^{t(k-j)}$$

「はい。続いて、見やすいように $l=k-j$ として、内側の総和を $l$ に関する総和に書き換えます」

$$=\sum_{j=0}^{N-1}a[j]W_N^{tj}\sum_{l=0}^{N-j-1} b[l] W_N^{tl}$$

「うーん、後ろの総和が $l=N-1$ まで回ってくれれば、$b$ の DFT になるんだけど、途中で止まってるね……」

「そうなんですよ。でも実は、こうしていいんです」

$$=\sum_{j=0}^{N-1}a[j]W_N^{tj}\sum_{l=0}^{N-1} b[l] W_N^{tl}$$

「え、いやいや、勝手に足しちゃダメじゃないの? 足した部分は 0 とは限らないんだし」

勝手に足したもの?:$\displaystyle \sum_{l=N-j}^{N-1} b[l] W_N^{tl}$

「違いますよ、先輩。外側のシグマもあるんですから、勝手に足したのはこれです」

勝手に足したもの:$\displaystyle \sum_{j=0}^{N-1}a[j]W_N^{tj}\sum_{l=N-j}^{N-1} b[l] W_N^{tl}$

「なるほど。でも結局は同じ話じゃない?」

「このとき、$l\ge N-j$ 、つまり $j+l\ge N$ であることに注目します。$N=N_a+N_b-1$ ですから、$j+l\ge N_a+N_b-1$ となります」

「ということは…… あー! なるほど、もし $j\le N_a-1$ かつ $l\le N_b-1$ だったら $j+l\le N_a+N_b-2$ になるのか。だから、勝手に足した部分については、絶対に $j\ge N_a$ または $l\ge N_b$ が成り立つんだね」

「そうです! したがって、勝手に足した部分について、必ず $a[j]=0$ か $b[l]=0$ のどちらかが成り立ち、$a[j]b[l]=0$ になります。だから、勝手に足した部分は必ず 0 です」

「すごい変形だね」

「結局、以下の式が成立します」

$$C[t]=\sum_{j=0}^{N-1}a[j]W_N^{tj}\sum_{l=0}^{N-1} b[l] W_N^{tl}$$

「右辺は、まず後半の総和が $j$ の値によらず、 $b$ の DFT 結果の $B[t]$ になっているね。そして、前半の総和が $a$ の DFT 結果の $A[t]$ になっているから、こうなるのか」

$$C[t]=A[t]B[t]$$

「はい。したがって、数列 $a,b$ の DFT 結果から数列 $c$ の DFT 結果を計算することは $O(N)$ でできます」

「畳み込みが、ただの要素ごとの積になるのか。面白いね。畳み込み計算についてまとめると、こういう手順だね。ここで、数列 $a,b$ の長さをそれぞれ $N_a,N_b$ としているよ」

  1. 数列 $a,b$ の長さが $N=N_a+N_b-1$ になるようにそれぞれ末尾に $0$ を追加
  2. 数列 $a,b$ を離散フーリエ変換し、数列 $A,B$ を得る
  3. 数列 $C$ を、$C[t]=A[t]B[t]$ で計算する
  4. 数列 $C$ を逆離散フーリエ変換し、数列 $c$ を得る。これが畳み込み結果である

「そういうことです。ただし、DFT や IDFT に $O(N^2)$ かかってしまうと普通に畳み込みするのと変わりません。そこで登場するのが高速フーリエ変換です!」

高速フーリエ変換(FFT)

高速フーリエ変換(FFT: Fast Fourier Transform)は、長さが $2^n$ である数列に対して、DFT と IDFT を $O(N\log N)$ で計算するアルゴリズムです」

「長さが 2 の累乗じゃないとダメなのか」

「はい。そうでない場合は、長さが 2 の累乗になるように末尾に $0$ を追加すればいいだけなので、競プロで使う上では影響のない制約ですけどね」

「で、どうやって計算するの?」

「まず、話を簡単にするため、$N=4$ としましょう。このとき、元の数列 $f[x]$ と変換後の数列 $F[t]$ は、行列を用いて書くとこのようになります」

$$\begin{pmatrix} F[0] \\ F[1] \\ F[2] \\ F[3] \end{pmatrix}=\begin{pmatrix} W_4^0 & W_4^0 & W_4^0 & W_4^0 \\ W_4^0 & W_4^1 & W_4^2 & W_4^3 \\ W_4^0 & W_4^2 & W_4^4 & W_4^6 \\ W_4^0 & W_4^3 & W_4^6 & W_4^9 \end{pmatrix}\begin{pmatrix} f[0] \\ f[1] \\ f[2] \\ f[3] \end{pmatrix}$$

「そうだね」

「ここで、右辺の $f[x]$ について、偶数要素を上に、奇数要素を下に配置し直します。つまり、上から 0, 2, 1, 3 という順番にします。行列の中身もこれに合わせて並べ替えます」

$$\begin{pmatrix} F[0] \\ F[1] \\ F[2] \\ F[3] \end{pmatrix}=\begin{pmatrix} W_4^0 & W_4^0 & W_4^0 & W_4^0 \\ W_4^0 & W_4^2 & W_4^1 & W_4^3 \\ W_4^0 & W_4^4 & W_4^2 & W_4^6 \\ W_4^0 & W_4^6 & W_4^3 & W_4^9 \end{pmatrix}\begin{pmatrix} f[0] \\ f[2] \\ f[1] \\ f[3] \end{pmatrix}$$

「なるほど、こうすると左 2 列について、$W_4^{2n}$ しか登場しないから、$W_2=W_4^2$ を使って書き直せるんだね」

「そのとおりです。勘がいいですね。また、$W_N^{nN+k}=W_N^k$ を利用して、こう書くことができます」

$$\begin{pmatrix} F[0] \\ F[1] \\ F[2] \\ F[3] \end{pmatrix}=\begin{pmatrix} W_2^0 & W_2^0 & W_4^0W_2^0 & W_4^0W_2^0 \\ W_2^0 & W_2^1 & W_4^1W_2^0 & W_4^1W_2^1 \\ W_2^0 & W_2^0 & W_4^2W_2^0 & W_4^2W_2^0 \\ W_2^0 & W_2^1 & W_4^3W_2^0 & W_4^3W_2^1 \end{pmatrix}\begin{pmatrix} f[0] \\ f[2] \\ f[1] \\ f[3] \end{pmatrix}$$

「うーん、ややこしいな。えっと、$W_2$ の要素に注目すると、次の行列が基本パターンになっていて、これが 4 ブロックに並べられているね」

$$\begin{pmatrix} W_2^0 & W_2^0 \\ W_2^0 & W_2^1 \end{pmatrix}$$

「はい。実はこれは、$N=2$ のときの DFT 行列なんです」

「あ、本当だ。この基本パターンが並んだ形にしたいから、右 2 列にある余分な項を $W_4^n$ の形で分離したんだね。で、$W_4$ の項にも規則性があるね」

「基本パターンを用いて、この行列を作ってみたいと思います。まず、左上のブロックと左下のブロックは、基本パターンのままなので、次の式で実現されます」

$$\begin{pmatrix} W_2^0 & W_2^0 \\ W_2^0 & W_2^1 \end{pmatrix}=\begin{pmatrix} 1 & 0 \\ 0 & 1 \end{pmatrix}\begin{pmatrix} W_2^0 & W_2^0 \\ W_2^0 & W_2^1 \end{pmatrix}$$

「うん。右上のパターンは、上段に $W_4^0$ が、下段に $W_4^1$ が掛かっているから、こう書けるね」

$$\begin{pmatrix} W_4^0W_2^0 & W_4^0W_2^0 \\ W_4^1W_2^0 & W_4^1W_2^1 \end{pmatrix}=\begin{pmatrix} W_4^0 & 0 \\ 0 & W_4^1 \end{pmatrix}\begin{pmatrix} W_2^0 & W_2^0 \\ W_2^0 & W_2^1 \end{pmatrix}$$

「はい。同様に、右下のパターンはこうなります」

$$\begin{pmatrix} W_4^2W_2^0 & W_4^3W_2^0 \\ W_4^2W_2^0 & W_4^3W_2^1 \end{pmatrix}=\begin{pmatrix} W_4^2 & 0 \\ 0 & W_4^3 \end{pmatrix}\begin{pmatrix} W_2^0 & W_2^0 \\ W_2^0 & W_2^1 \end{pmatrix}$$

「これが 4 つ並ぶように行列をつくればいいから、こうなるね」

$$\begin{pmatrix} F[0] \\ F[1] \\ F[2] \\ F[3] \end{pmatrix}=\begin{pmatrix} 1 & 0 & W_4^0 & 0 \\ 0 & 1 & 0 & W_4^1 \\ 1 & 0 & W_4^2 & 0 \\ 0 & 1 & 0 & W_4^3 \end{pmatrix} \begin{pmatrix} W_2^0 & W_2^0 & 0 & 0 \\ W_2^0 & W_2^1 & 0 & 0 \\ 0 & 0 & W_2^0 & W_2^0 \\ 0 & 0 & W_2^0 & W_2^1 \end{pmatrix}\begin{pmatrix} f[0] \\ f[2] \\ f[1] \\ f[3] \end{pmatrix}$$

「こうなったけど、これをどう解釈すればいいんだろう」

「まず、右側の行列を掛ける部分に注目します」

$$\begin{pmatrix} W_2^0 & W_2^0 & 0 & 0 \\ W_2^0 & W_2^1 & 0 & 0 \\ 0 & 0 & W_2^0 & W_2^0 \\ 0 & 0 & W_2^0 & W_2^1 \end{pmatrix}\begin{pmatrix} f[0] \\ f[2] \\ f[1] \\ f[3] \end{pmatrix}$$

「この結果は、もとの数列の偶数要素 $f[0],f[2]$ を $N=2$ で DFT した結果と、奇数要素 $f[1],f[3]$ を $N=2$ で DFT した結果が並ぶね」

「はい。並べ替えた数列で考えると、前半の DFT 結果と後半の DFT 結果が並びますね。この、$N=2$ での DFT 結果を $F_2[0],F_2[1],F_2[2],F_2[3]$ とすると、こうなります」

$$\begin{pmatrix} F_2[0] \\ F_2[1] \end{pmatrix}=\begin{pmatrix} W_2^0 & W_2^0 \\ W_2^0 & W_2^1 \end{pmatrix}\begin{pmatrix} f[0] \\ f[2] \\ \end{pmatrix}$$

$$\begin{pmatrix} F_2[2] \\ F_2[3] \end{pmatrix}=\begin{pmatrix} W_2^0 & W_2^0 \\ W_2^0 & W_2^1 \end{pmatrix}\begin{pmatrix} f[1] \\ f[3] \\ \end{pmatrix}$$

「したがって、最終的な数列 $F$ の計算は、この計算になります」

$$\begin{pmatrix} F[0] \\ F[1] \\ F[2] \\ F[3] \end{pmatrix}=\begin{pmatrix} 1 & 0 & W_4^0 & 0 \\ 0 & 1 & 0 & W_4^1 \\ 1 & 0 & W_4^2 & 0 \\ 0 & 1 & 0 & W_4^3 \end{pmatrix}\begin{pmatrix} F_2[0] \\ F_2[1] \\ F_2[2] \\ F_2[3] \end{pmatrix}$$

「この行列の掛け算はシンプルだね。つまり、$F[t]$ はそれぞれこうして得られるということだよね」

$$F[0]=F_2[0]+F_2[2]W_4^0$$

$$F[1]=F_2[1]+F_2[3]W_4^1$$

$$F[2]=F_2[0]+F_2[2]W_4^2$$

$$F[3]=F_2[1]+F_2[3]W_4^3$$

「はい。また、$W_4^2=-1$ を利用すると、こう書けます」

$$F[0]=F_2[0]+F_2[2]W_4^0$$

$$F[1]=F_2[1]+F_2[3]W_4^1$$

$$F[2]=F_2[0]-F_2[2]W_4^0$$

$$F[3]=F_2[1]-F_2[3]W_4^1$$

「こうすると、$F_2[0],F_2[1]$ と $F_2[2]W_4^0,F_2[3]W_4^1$ しか出てこなくなるね」

「はい。これを図に書くと、こうなります」

「青い矢印はそのまま、赤い矢印は書かれている係数を掛けてから足すということだね」

「はい。この矢印の形をよく見ると、8 の字の形が 2 つ重なっているように見えるので、この形を蝶になぞらえて “バタフライ演算” と呼ばれます」

「これで、$N=2$ での DFT が計算できていれば、$N=4$ の DFT はバタフライ演算を用いて $O(N)$ で解ける、ということが分かったけど、$N=8$ の場合はどうなるの?」

「その前に、$N=2$ の場合についてもう一度考えてみましょう」

「$N=2$ って、こうでしょ?」

$$\begin{pmatrix} F_2[0] \\ F_2[1] \end{pmatrix}=\begin{pmatrix} W_2^0 & W_2^0 \\ W_2^0 & W_2^1 \end{pmatrix}\begin{pmatrix} f[0] \\ f[2] \\ \end{pmatrix}$$

$$\begin{pmatrix} F_2[2] \\ F_2[3] \end{pmatrix}=\begin{pmatrix} W_2^0 & W_2^0 \\ W_2^0 & W_2^1 \end{pmatrix}\begin{pmatrix} f[1] \\ f[3] \\ \end{pmatrix}$$

「はい。これも、次のように解釈することで、バタフライ演算で表現できます

$$\begin{pmatrix} F_2[0] \\ F_2[1] \end{pmatrix}=\begin{pmatrix} 1 & W_2^0 \\ 1 & -W_2^0 \end{pmatrix}\begin{pmatrix} f[0] \\ f[2] \\ \end{pmatrix}$$

$$\begin{pmatrix} F_2[2] \\ F_2[3] \end{pmatrix}=\begin{pmatrix} 1 & W_2^0 \\ 1 & -W_2^0 \end{pmatrix}\begin{pmatrix} f[1] \\ f[3] \\ \end{pmatrix}$$

「なるほど、つまり、$N=4$ の DFT は、並べ替え→ $N=2$ のバタフライ演算(2 ブロック)→ $N=4$ のバタフライ演算、の 3 ステップで計算できるんだね。」

「ということは、$N=8$ の場合は、さらにバタフライ演算が 1 段増えるということ?」

「はい。そうなります。ただし、最初の並べ替えに注意しないといけません」

「最初の並べ替えって、偶数要素を最初に持ってくる (0, 2, 4, 6, 1, 3, 5, 7) じゃないの?」

「それは、$N=8$ のバタフライ演算をする直前の状態です。$N=4$ のバタフライ演算後にその並びになっている必要があります。前半は、(0, 2, 4, 6) の偶数要素 0, 4 を前に持ってきた (0, 4, 2, 6) になっていないといけません。後半も同様に (1, 5, 3, 7) です」

「ややこしいね。つまり、最初に入力配列を (f[0], f[4], f[2], f[6], f[1], f[5], f[3], f[7]) に並べ替えてから、$N=2,4,8$ のバタフライ演算を順番に行うと、正しい並びになるってことか」

「はい。この (0, 4, 2, 6, 1, 5, 3, 7) という並びは、一見奇妙な並び方ですけど、実は綺麗な法則があるんです。それは、(0, 1, 2, 3, 4, 5, 6, 7) を二進数で表現してひっくり返した値になっているという法則です!」

「二進数でひっくり返す?」

「1 は二進数 3 桁表現で 001 です。これをひっくり返すと、001 となり、 4 です。他のも、ほら!」

「おお、すごいね。本当に綺麗な対応だ」

「その並べ替えを終えたら、あとはバタフライ演算を $N=2,4,8$ で繰り返すだけです」

「$N=8$ のときのバタフライ演算はどうなるの?」

「こういう図になります。図を書くのは $N=8$ が限界ですね……」

「規則的だね。並べ替えに $O(N)$、バタフライ演算が $O(\log N)$ 段あって、それぞれの段で $O(N)$ で計算できるから、全体で $O(N\log N)$ というわけか」

「はい。じゃあ、このアルゴリズムを書いていきましょう!」

Cooley–Tukey 型 FFT アルゴリズムの実装

「今考えた、並べ替えとバタフライ演算で DFT 計算をするアルゴリズムを、Cooley-Tukey 型 FFT アルゴリズムと言います。Cooley と Tukey は考案者の名前ですね。このアルゴリズムが FFT の中では最もよく使われます」

「へえ、他にもあるのか」

「はい。Cooley-Tukey 型は、in-place でできる、つまり余計な配列を必要としないのが特徴です」

「じゃあ、考えていこうか。まず、並べ替えだね」

バタフライ演算用の並べ替えは、これでできます」

「中の for ループに複雑な計算があるけど、これは何をしているの?」

「これは、$i$ の第 $k$ 桁の値を、$j$ の第 $h-1-k$ 桁に設定しています。配列の長さを $n=2^h$ としています。添字の桁数は $h$ 桁です」

「なるほど、(i>>k&1) で $i$ の第 $k$ 桁を取ってきて、それをビットシフトで第 $h-1-k$ 桁にして、$j$ との bitwise-or を取っているんだね」

「はい。そのままひっくり返しているだけですね」

「そして、$a[i]$ と $a[j]$ を入れ替えているんだね。ここで、 2 回入れ替えると意味がないから、$i<j$ のときのみ入れ替えるようにしているんだね」

「続いて、バタフライ演算です」

「最初のステップを第 $0$ ステップとすると、第 $s$ ステップでは、長さ $2^{s+1}$ のブロックがあるね」

「はい。ここで、段数でループを回すのではなく、ブロックサイズの半分でループを回しましょう」

「ブロックサイズの半分? どうして?」

「ブロックサイズの半分の値は使いますが、段数の値は使用しないので。また、ブロックサイズでループすると 2 で割る必要が出てくるので、半分を使います」

「へー、そういうもんなのか」

「ブロックサイズの半分を $b$ とすると、最初は $b=1$ で、段が進むごとに 2 倍になっていきます。そして、$b=n$ になると終わりです。なので、ループはこうなりますね」

「さて、各ステップでの処理を考えよう。図がないと考えにくいね」

「ここのループのコツがあって、ブロックをループしていくのではなく、各ブロック内で $0,0+b$ 番目の値を全部処理してから、各ブロック内で $1,1+b$ 番目の値を処理する、という形にループしていきます」

「なるほど、$j,j+b$ 番目の値の処理ではどのブロックでも、$j+b$ 番目の要素に掛ける重み係数が $W_{2b}^{j}$ だから、係数を使い回せるんだね」

「はい。ループはこうなります」

complex<double> というのは、実部と虚部が double 型で表現される複素数だね。そして std::polar(r, t) は、絶対値 r 、偏角 t の複素数 $re^it$ を生成する関数だね。どっちも #include <complex> で使えるかな。今欲しいのは $W_{2b}^j$ だから、$e^{-i\frac{2\pi j}{2b}}$ だね」

「はい。最後に、具体的な計算を書いていきます。ブロックの先頭を $k$ として、$k=0,2b,4b,\dots$ とループしていきます」

「そうすると、考えるのは、$a[j+k],a[j+k+b]$ の値から新たな $a[j+k],a[j+k+b]$ を計算することだね」

「することは単順ですね。$s=a[j+k],t=a[j+k+b]\times W_{2b}^j$ とし、$a[j+k]\leftarrow s+t,a[j+k+b]\leftarrow s-t$ とします」

「さて、これで終わり?」

「はい。終わりです。FFT 関数の中身はこんな感じですね。入力の数列 $a$ は、複素数配列 vector<complex<double>> とします」

「これだけか。思ったより簡単だね」

「はい。続いて、 IFFT をこの中に組み込みましょう」

「そうか、逆変換があったね」

「といってもやることは変わらなくて、重み係数を $W_{2b}^{-j}$ にするのと、最後に $n$ で割るだけです。なので、fft() の引数に bool inverse (デフォルトで false)を追加して、条件分岐するようにしましょう。こんな感じです」

「本当に少しの変更で逆変換もできるんだね」

「はい。実際に使うときは vector<double> が入力される場合が多いので、そっちにも対応しておきましょう」

「関数をオーバーロードしているんだね」

「では、続いて、これを用いた畳み込みの実装をしてみましょう」

FFT を用いた畳み込みの実装

「これに関しては、実装を見てくださいという感じですね」

「えっと、畳み込み結果の長さが a.size()+b.size()-1 になるから、これ以上の 2 の累乗数を t として計算しているんだね。そして、a, b をリサイズして FFT し、新たな複素数列 A, B を得ているね」

「そして、その積を計算します。新たな数列を用意するのが無駄なので、A を使いまわしています。A[i] *= B[i]; が全て終われば、A には畳み込み結果を FFT したものが入っています」

「最後に、得られた複素数列を IFFT して畳み込み結果を得るんだね。ここでも A を使いまわしているね。IFFT 結果は複素数列として得られるけど、元の a, b が実数列だから、結果は実数列になるはずだ。だから、実部を取り出して答えになると」

「はい。畳み込み結果の本来の長さは s なので、出力する前に a を長さ s にリサイズしておかないの変なことになるので要注意ですね」

AtCoder Typical Contest 001 C「高速フーリエ変換」

「さて、コードが合っているかどうか、AtCoder で提出して verify したいな」

「AtCoder には、FFT の verify 用の問題、『高速フーリエ変換』がありますよ」

「典型 verify 用コンテストだね。じゃあこれに提出してみよう」

「この問題では、a[1]〜a[N] 、b[1]〜b[N] が与えられるので、添字の和が k になるような値のペアの積の和を、k=1〜2N まで出力します」

「そのまんま畳み込みだね」

「配列 a, b のサイズを N+1 として、a[0]=b[0]=0 としておくと分かりやすいですね」

「そうだね。入力を配列に入れたら、畳み込んで、出力するだけだね」

「出力するときに、0 番目を出力しないように気をつけましょう」

「ということで、解答コードの main 関数はこうかな」

「サンプルは合いましたね」

「この問題では、N が 10^5 以下、a[i], b[i] が 10^2 以下だから、出力する値は 10^9 以下だね。だから、最大ケースでも int に収まるし、途中で double 型で処理しても、精度は大丈夫だね」

「double だと 10 進数で 15 桁程度の精度なので、 10^18 くらいの値の出力だと long double を使うなどしないといけませんね」

「そうだね。じゃあ、提出……と。お、AC だね」

「今回の問題では FFT で大丈夫でしたけど、もっと値が大きくなって、M で割ったあまりを出力する問題になったらダメですね」

「そうだね。その場合、確か特定の素数 mod 上で FFT をする、数論変換(NTT)というのがあるんだよね」

「はい。せっかく FFT を勉強したんですから、こっちもやってみたいですね。私も NTT は使ったことないんですけど」

「NTT はすごく高度なものというイメージがあるけど、挑戦してみたいね」

提出コード

Pocket

スポンサーリンク
レクタングル大
レクタングル大

シェアする

  • このエントリーをはてなブックマークに追加

フォローする