はじめに

競技プログラミング(AtCoder)をしていると、水レートあたりから顔を出し始めて、青になってくるとみんな知ってる、そんなアルゴリズム、畳み込み。僕もようやくこれに手を付けて、少しばかり勉強して、自分のライブラリに入れました。

畳み込み自体は数年前から知っていたのですが、複素数やら、整数環やらと出てくるとかなり構えてしまい、また読み進めると何をやっているのかよくわからなくなって、幾度も理解を諦めていました。今回の記事では、よくある他の記事とは順序を変えて、僕が理解した順序で説明を書いてみます。複雑な数学的背景は可能な限り飛ばしますが、どうしても数式を使うことは避けられないので、その点は頑張ってください。また、具体的なコードの実装の説明も行いません。あくまでも理論的な部分を理解する土台を作るための記事です。

またtatyamさんが、前にも畳み込みに使われるアルゴリズムであるFFTの解説を書いて下さっています。この記事よりも正確になっているので、数式が全く怖くない、FFTの仕組みをしっかり理解したい、という方はこちらを見た方がいいかもしれません。

【競プロer向け】FFT を習得しよう！

こんにちは、20B の @tatyam [https://trap.jp/author/tatyam] です。この記事は、traP夏のブログリレー[/tag/summer-blog-relay] 33日目の記事です。 高速フーリエ変換 といえば、競技プログラミングで数列の畳み込みなどに使われるテクニックですね。今日は高速フーリエ変換と高速フーリエ変換を利用した畳み込みを習得していきましょう。 ポエム一般的なフーリエ変換の 定義 / 説明 と異なる部分があるかもしれませんが、競プロであればこれで十分だと思います。AtCoder では AC Library[https://atcod…

東京科学大学デジタル創作同好会traPtatyam

畳み込みとは

定義

原義的な畳み込み、特に離散的な方の畳み込みは、次のようにして定義されます。

$$(f\ast g)(m)=\sum _{n}f(n)g(m-n)$$

言葉で言うと、二つの関数の引数の和が一定になるようにして足し合わせたものです。ここで出てきた関数 $f$ や $g$ は、この記事では数列と解釈して差し支えないでしょう。そこで、この畳み込みを二つの数列$A$と$B$に対して定義して、その結果を$C$とします。すると、次のように書くことができます。

$$C_n=\sum _{m=0}^n{A}_m{B}_{n-m}$$

競技プログラミングで出てくる畳み込みは、この式で計算される $C_n$ を、定義できる全ての$n$について求める計算を指します。

計算量

この畳み込みを計算することを考えます。各$n$について$n$個の総和を計算することになるので、時間計算量は$O(N^2)$になります。これでは遅いので、もう少し早くしたいところです。線形時間とは言いませんが、$O(N\log N)$程度になってくれると嬉しい。

そしてこの計算量は実現可能で、これを実現するのが 高速フーリエ変換　(FFT) と呼ばれるアルゴリズムです。ここからは、畳み込みがどう高速化されているのか、FFTは何に使われるのかをゆっくりと理解していきましょう。

数列から多項式へ

畳み込みを理解する前に、少し数列の表現方法について説明します。

係数による表現

フーリエ変換を用いた畳み込みの計算では、数列を多項式として解釈します。つまり、長さ$N$の数列$A=(a_0,a_1,...,a_{n-1})$があるとき、これを多項式$f_A(x)=a_0+a_{1}x+a_2{x}^2+...+a_{n-1}{x}^{n-1}$で表される$N-1$次の多項式として見ます。こうすると、畳み込みを多項式の積として解釈できます。

例として長さが3と2の数列$A=(a_0,a_1,a_2)$と$B=(b_0,b_1)$を用意し、それぞれに対応する多項式を$f_A(x)$、$f_B(x)$とします。この積$f_C(x)$を計算してみましょう。

$$\begin{gathered} f_C(x)=f_A(x)\times f_B(x)=(a_0+a_{1}x+a_2{x}^2)(b_0+b_{1}x) \\ =(a_0{b}_0)+(a_0{b}_1+a_1{b}_0)x+(a_1{b}_1+a_2{b}_0)x^2+(a_2{b}_1)x^3 \end{gathered}$$

各項の係数についている添え字に注目します。すると、どの項でも$A$の添え字と$B$の添え字の和がその項の$x$の次数に一致することがわかります(2乗の項の$a_0{b}_2$は$b_2=0$と解釈する影響で消えてしまっています)。特に、$x$の$n$乗の項の係数を$c_n$として$C=(c_0,c_1,c_2,c_3)$とすると、$C$は$A$と$B$の畳み込みになっています。

値による表現

前項では数列から多項式を作ることができることがわかりました。今度は、一つの多項式を数列で表現することを考えます。各行の係数を取ってきた数列は、元の多項式をよく表現してくれます。ところで、多項式を表現する、一つに指定する方法はそれだけでしょうか？

直線は1次の多項式を使って表すことができます。そして、その直線を特定するには、その直線が通る2点の情報があれば十分です。2次関数なら2次の多項式を用いて、通過する3点があればよく、3次関数では3次の多項式で4つの通過点があると特定できます。

一般に、$n$次の多項式$f(x)=a_0+a_{1}x+...+a_n{x}^n$の各係数は、グラフ$y=f(x)$の通過する点が$n+1$個あれば求めることができます。言い換えると$n+1$個(以上)の点の情報$(f(r_0),f(r_1),...,f(r_n))$があれば、一つの多項式を表現できるということです。

尚、各多項式がどの数列によるものなのかをわかりやすくするため、数列$A$から作られた多項式$f_A$を単に$A(x)$、その出力値から得られる数列を$\hat{A}=({\hat{a}}_0,{\hat{a}}_1,...,{\hat{a}}_{n-1})$とします。
(数列$A$と混同しないよう気を付けてください)。

多項式の掛け算

さて、数列から作られる多項式の表現として

1. 係数からなる数列
2. 多項式の出力値からなる数列

の2種類を得ることができました。今回やりたいことである畳み込みは多項式の積と対応するので、多項式の積を取ったときにこれらの数列がどう変わるのかを調べてみます。これまで通り、数列$A$と$B$からなる多項式$A(x)$と$B(x)$を考えます。

まず、係数からなる数列は先述の通り二つの数列の畳み込みとなります。問題は出力値の方です。こちらは単なる値の積となってくれます。つまり、$C(x)$について、

$$(C(r_0),C(r_1),...,C(r_{n-1}))=(A(r_0)B(r_0),A(r_1)B(r_1),...,A(r_{n-1})B(r_{n-1}))$$

です。$C(x)=A(x)B(x)$で定義したので当然ですが、こちらはかなり簡単になりました。$\hat{A}$と$\hat{B}$さえわかれば$\hat{C}$の値が即座にわかります。

問題の変換

$C(r_i)=A(r_i)B(r_i)$と、積の計算が随分と楽になりました。しかし、これで終わりではありません。ここまで言及していなかった、重要な事項がいくつかあります。

まずは$A(r_i)$を実際に計算するときの計算量です。$A(x)$は$N-1$次の多項式ですから、この値を愚直に計算するには$O(N)$の時間計算量が必要になります。これを必要なだけ求めると、$\hat{A}$の計算には最終的に$O(N^2)$時間必要となり、高速化は実現されていません。
第二に、$C$から$\hat{C}$への変換に使う引数$r_i$の決定です。どのような値を使えばよいでしょうか？
第三に、$\hat{C}$から$C$への変換です。ここまで、$C$から$\hat{C}$の変換方法は説明しましたが、その逆はまだです。

次の章からは、数列の畳み込みを多項式の積と解釈したことによって出たこの課題を解決する手段を説明します。少しフライングすると、$r_i$を適切に決めることで、これらの問題を一気に解決します。

高速フーリエ変換

入力 $r_i$ の決定

$N$要素の数列$A$から$\hat{A}$を計算することで、畳み込みの計算が1歩進むことがわかりました。ここからはその計算方法を実際に追っていきます。

まずは$A(x)$への入力とする$r_i$の決定です。$r_i$はどんな値がいいでしょうか？まず、異なる$N$個の点が欲しいので、$r_i$も互いに異なっていてほしいですね。他にどんな条件があるといいでしょうか。離散フーリエ変換で用いられる$r_i$は、追加で次の性質を持っています。

• $r_i$は次の性質を満たす数$\omega$の$i$乗
  • ${\omega }^N=1$
  • ${\omega }^{(N/2)}=-1$ ($N$が偶数のとき)

アルゴリズムの解説に集中したいので、この記事では具体的な$\omega$の値は説明しません。ほとんどのFFTの解説記事には書いてあるので、気になる人は調べに行ってみましょう。

$A(r_i)$の計算

では、このような値を使って計算したとき、どんなことが起こるでしょうか。
ここでは、$N$を偶数として計算をしてみます。(もし$N$が偶数じゃないなら、数列に0の項を一つ付け足せば良いです)
まず、$A(x)$は${\sum }_{j=0}^{N-1}{a}_j{x}^j$ ですから、

$${\hat{a}}_i=\sum _{j=0}^{N-1}{a}_j{r}_i^j=\sum _{j=0}^{N-1}{a}_j{\omega }^{ij}$$

と書けます。$N$が偶数なので、この総和について$j$が偶数の項と奇数の項とで分けてみます。

$${\hat{a}}_i=\sum _{j=0}^{N-1}{a}_j{\omega }^{ij}=\sum _{j=0}^{N/2-1}{a}_{2j}{\omega }^{2ij}+\sum _{j=0}^{N/2-1}{a}_{2j+1}{\omega }^{i(2j+1)}$$

ここの第二項について、${\omega }^i$をくくって

$${\hat{a}}_i=\sum _{j=0}^{N/2-1}{a}_{2j}{\omega }^{2ij}+{\omega }^i\sum _{j=0}^{N/2-1}{a}_{2j+1}{\omega }^{2ij}$$

となります。では同じ計算を、${\hat{a}}_{i+N/2}$についても行ってみると、なんやかんやあって ($i$を$i+N/2$に変えるだけです)

$${\hat{a}}_{i+N/2}=\sum _{j=0}^{N/2-1}{a}_{2j}{\omega }^{2(i+N/2)j}+{\omega }^{i+N/2}\sum _{j=0}^{N/2-1}{a}_{2j+1}{\omega }^{2(i+N/2)j}$$

となります。$\omega$周りは、$\omega$の持つ性質、${\omega }^N=1$、${\omega }^{(N/2)}=-1$を使ってさらに整理できて

$${\hat{a}}_{i+N/2}=\sum _{j=0}^{N/2-1}{a}_{2j}{\omega }^{2ij}-{\omega }^i\sum _{j=0}^{N/2-1}{a}_{2j+1}{\omega }^{2ij}$$

となります。

この二つを見比べてみると、なんと足し引きが違うだけでほとんど同じであることがわかります。試しに

$$\begin{gathered} {\hat{e}}_i=\sum _{j=0}^{N/2-1}{a}_{2j}{\omega }^{2ij} \\ {\hat{o}}_i=\sum _{j=0}^{N/2-1}{a}_{2j+1}{\omega }^{2ij} \end{gathered}$$

なんて書けば、

$$\begin{gathered} {\hat{a}}_i={\hat{e}}_i+{\omega }^i{\hat{o}}_i \\ {\hat{a}}_{i+N/2}={\hat{e}}_i-{\omega }^i{\hat{o}}_i \end{gathered}$$

となります。つまり、実質的に計算が必要な${\hat{a}}_i$が半分になっています。
そして、${\hat{e}}_i$や${\hat{o}}_i$の式を見ると、これもまた同じ出力値の計算になっていることがわかります。

再帰計算

少し詳しく書きます。$E=(e_0,e_1,...,e_{N/2})=(a_0,a_2,a_4,...,a_{n-1})$ として、$\hat{E}$を考えます。$E$は$A$のちょうど半分から作っているので、その要素数が$N/2$個であることに注意すると、計算に使う$r_i$として$({\omega }^2{)}^i$を使うことができます。実際、

• $({\omega }^2{)}^{N/2}={\omega }^N=1$
• $({\omega }^2{)}^{(N/2)/2}={\omega }^{N/2}=-1$

となって、$A$に対する計算で使った$\omega$と同様の性質を満たします。

では計算してみましょう。$E(r_i)$を求めます。

$$\begin{gathered} E(r_i)=\sum _{j=0}^{N/2-1}{e}_j{r}_i^j=\sum _{j=0}^{N/2-1}{a}_{2j}({\omega }^2{)}^{ij} \\ ={\hat{e}}_i \end{gathered}$$

となって、確かに${\hat{e}}_i$が$E(r_i)$となっています。同様に、${\hat{o}}_i=O(r_i)$です。

ここまでをまとめると、$\hat{A}$($N$要素)を求める問題を、$\hat{E}$と$\hat{O}$ ($N/2$要素)を求める問題に分割することができました。もし$E$や$O$も偶数個からなる場合、これをさらに分割できます。もし$N$が2冪であれば、この分割を繰り返すことができて、高速に計算できそうです。特に、数列の末尾に0をつけ足しても結果は変わらないので、$N$はいつでも2冪にすることができます。ﾔｯﾀｰ

こんな風に偶奇で分割して計算していくのをバタフライ演算とか呼んだりするらしいです。僕にはどこが蝶なのかわからなかった...

これによる計算量を見積もります。$A$を$E$と$O$に分割して同じ問題を解いた後、$\hat{A}$はその結果$\hat{E}$と$\hat{O}$から各要素を定数時間で計算できる(${\hat{a}}_i={\hat{e}}_i+{\omega }^i{\hat{o}}_i,{\hat{a}}_{i+N/2}={\hat{e}}_i-{\omega }^i{\hat{o}}_i$)ので、後半の計算は$O(N)$となります。そして、数列の分割は$O(\log N)$段階まで出来るので、$A$から$\hat{A}$の計算は全体で$O(N\log N)$時間で計算できました。

尚、計算を二つに分け、分割しながら計算したこのアルゴリズムを高速フーリエ変換 (FFT) と呼びます。

逆変換

ここで目的を再確認しましょう。当初の目的は数列$A$と$B$の畳み込み$C$を計算することです。そのため、$A$や$B$で作られる多項式$A(x)$と$B(x)$の積を考えました。ここで、$A(x)$や$B(x)$にいくつかの値を代入して得られる数列$\hat{A}$と$\hat{B}$を使うと、$C(x)$を表現している$\hat{C}$を簡単に求めることができます。

ここまで、そのうち$A$、もとい$A(x)$から$\hat{A}$を高速に計算する方法を解説しました。あとは$\hat{C}$から$C(x)$を復元し、$C$を得ることができれば目的達成です。

結論から書くと、この逆変換は

$$a_i=\frac{1}{N}\sum _{j=0}^{N-1}{\hat{a}}_j{r}_i^{-j}=\frac{1}{N}\sum _{j=0}^{N-1}{\hat{a}}_j{\omega }^{-ij}$$

です。
これが正しいことを実際に確かめます。まずは左辺に${\hat{a}}_j$の計算式を代入します。

$$\frac{1}{N}\sum _{j=0}^{N-1}\left(\sum _{k=0}^{N-1}{a}_k{\omega }^{jk}\right){\omega }^{-ij}$$

和の順序を入れ替えて

$$\frac{1}{N}\sum _{k=0}^{N-1}{a}_k\sum _{j=0}^{N-1}{\omega }^{j(k-i)}$$

ここで出てきた${\sum }_{j=0}^{N-1}{\omega }^{j(k-i)}$をさらに計算します。
$i=k$のとき、${\omega }^{j(k-i)}=1$なので、${\sum }_{j=0}^{N-1}{\omega }^{j(k-i)}=N$です。

$i\ne k$なら、${\sum }_{j=0}^{N-1}{\omega }^{j(k-i)}$は初項$1$、公比${\omega }^{k-i}$の等比数列の和になります。(${\omega }^{k-i}$は$1$にならないという性質を持ちます。)

そのため、等比数列の和の公式から

$$\sum _{j=0}^{N-1}{\omega }^{j(k-i)}=\frac{({\omega }^{k-i}{)}^N-1}{{\omega }^{k-i}-1}=\frac{({\omega }^N{)}^{k-i}-1}{{\omega }^{k-i}-1}=\frac{1^{k-i}-1}{{\omega }^{k-i}-1}=0$$

です。以上まとめると、

$$\frac{1}{N}\sum _{k=0}^{N-1}{a}_k\sum _{j=0}^{N-1}{\omega }^{j(k-i)}$$

は$i=k$の場合だけが残って、

$$\frac{1}{N}\sum _{k=0}^{N-1}{a}_k\sum _{j=0}^{N-1}{\omega }^{j(k-i)}=\frac{1}{N}\times N{a}_i=a_i$$

となって、確かに逆変換が成立していることがわかりました。

順変換との差異

$a_i$と${\hat{a}}_i$の相互変換を求めることができました。
2つの式を比べてみます。

$$\begin{gathered} {\hat{a}}_i=\sum _{j=0}^{N-1}{a}_j{\omega }^{ij} \\ {a}_i=\frac{1}{N}\sum _{j=0}^{N-1}{\hat{a}}_j{\omega }^{-ij} \end{gathered}$$

まず逆変換の方は、和の後に$N$で割っています。もう一点、$\omega$が${\omega }^{-1}$に置き換わっていますね。
それ以外は順変換と逆変換は何も変わっていません。はい、何も変わっていません。つまり、順変換のアルゴリズムを逆変換にも使うことができ、順変換、逆変換ともに高速に計算出来るのです。

まとめ

さて、ここまでをまとめます。

二つの数列$A=(a_0,a_1,...,a_{N-1})$と$B=(b_0,b_1,...,b_{M-1})$について、その畳み込み$C=(c_0,c_1,...,c_{N+M-1})$を求めることを考えます。

まず、整数$L$を、$N+M-1$以上となる2の累乗で求めて、$A$と$B$の長さを$L$にします。(新しい部分は0で埋める)
次に$A$から

$${\hat{a}}_i=\sum _{j=0}^{L-1}{a}_j{\omega }^{ij}$$

で得られる$\hat{A}=({\hat{a}}_0,{\hat{a}}_1,...,{\hat{a}}_{L-1})$を高速フーリエ変換により計算し、同様に$\hat{B}$も求めます。これら二つから

$${\hat{c}}_i={\hat{a}}_i{\hat{b}}_i$$

で$\hat{C}=({\hat{c}}_0,{\hat{c}}_1,...,{\hat{c}}_{L-1})$が得られます。
これに対して、

$$c_i=\frac{1}{N}\sum _{j=0}^{L-1}{\hat{c}}_j{\omega }^{-ij}$$

で得られる$C$を再び高速フーリエ変換により計算することで、最初の目的を達成することができました。

以上のアルゴリズムが高速フーリエ変換により畳み込みを高速に計算する方法です。

最後に

(説明のまとめに「まとめに」を使ってしまったので「最後に」です)

お疲れさまでした！！！

🎉これであなたも畳み込み・高速フーリエ変換を理解できました！🎉

と言ってみたものの、実はまだ一つ巨大な問題が2つほど残っています。

まず一つ目は、今回の記事で出てきた謎の数$\omega$の正体です。個人的に、この$\omega$がFFTの理解における一番の壁だと思っています。水や青コーダーでこの類の分割統治につまずく人はほとんどいませんからね。今回、$\omega$は単なる都合のいい数としてしか出ていません。しかし、その性質だけを抑えておくことで、難しい部分をスキップして説明をすることができています。

二つ目は、この畳み込み、およびFFTを処理する実装です。この記事では理論的な部分の(それも一部の)説明にとどめています。今回説明しなかったこれら二つはまた別の記事としていずれ出す、かもしれません。しかし、ここまで理解出来たら、$\omega$の理解も、実装の理解も容易いはずです。

では、またの機会に。

参考

【競プロer向け】FFT を習得しよう！

こんにちは、20B の @tatyam [https://trap.jp/author/tatyam] です。この記事は、traP夏のブログリレー[/tag/summer-blog-relay] 33日目の記事です。 高速フーリエ変換 といえば、競技プログラミングで数列の畳み込みなどに使われるテクニックですね。今日は高速フーリエ変換と高速フーリエ変換を利用した畳み込みを習得していきましょう。 ポエム一般的なフーリエ変換の 定義 / 説明 と異なる部分があるかもしれませんが、競プロであればこれで十分だと思います。AtCoder では AC Library[https://atcod…

東京科学大学デジタル創作同好会traPtatyam