多項式の掛け算をFFTでやるやつを実装
この記事は「みゅーもり Advent Calendar 2018」の4日目分となります:
こんなにふざけたタイトルのAdvent Calendarが25日分すべて埋まるとは正直思っていませんでした. 登録してくださったみなさま, 本当にありがとうございます.
さて今日の内容は表題の通り, FFTを使って多項式同士の積の計算を高速化しようという話です.
専攻によっては学部の講義で扱う有名事実だそうですが, 無知なわたしはつい先日まで知りませんでしたし, これほど単純な問題に対して計算量を改善するテクニックがあるのかと驚いた話でもあったので, まとめておこうと思います.
Jupyter Notebook: http://nbviewer.jupyter.org/github/myuuuuun/blog_contents/blob/master/notebooks/polynomial_multiplication.ipynb
追記
- (2018/12/5)添字のミスを修正しました. @stmfj さんありがとうございます!
問題設定
を 次多項式
\begin{align} f(x) = a_0 + a_1 x + a_2 x^{2} + \ldots + a_{n - 1} x^{n - 1}, \\ g(x) = b_0 + b_1 x + b_2 x^{2} + \ldots + b_{n - 1} x^{n - 1} \end{align}
とする( )*1. このとき, 2つの多項式の積 は次多項式 \begin{align} h(x) = c_0 + c_1 x + c_2 x^{2} + \ldots + c_{2n-2} x^{2n-2} \end{align} となる(). この を求め, 多項式 を計算したい.
例
, のとき, は \begin{align} h(x) = 3x^{3} + 10 x^{2} + 17x + 12 \end{align} となる
素直な解法
単純に考えれば, の次の係数は \begin{align} c_k = \sum_{i=0}^{k} a_i b_{k-i} \end{align} で与えられるので, これをについて計算をすればよいでしょう. Pythonで実装すると
import numpy as np def get_poly_product(f_coefs, g_coefs): """ Parameters: f_coefs: array_like. fの係数のリスト. [0次の係数, 1次の係数, ..., 最高次の係数] g_coefs: array_like. gの係数のリスト. [0次の係数, 1次の係数, ..., 最高次の係数] Returns: h_coefs: ndarray. h=fgの係数のリスト. [0次の係数, 1次の係数, ..., 最高次の係数] """ f_size, g_size = len(f_coefs), len(g_coefs) h_coefs = np.zeros([f_size+g_size-1], dtype=float) for i in range(f_size): for j in range(g_size): h_coefs[i+j] += f_coefs[i] * g_coefs[j] return h_coefs
のようになります. 上の例題を解いてみると
get_poly_product([3, 2, 1], [4, 3]) >> array([12., 17., 10., 3.])
となり, 手で解いた結果と一致しました.
計算量
コードから明らかなように, この方法の計算量は になります. さて, これよりも少ない計算量で の係数を求めることはできるでしょうか?
実は以下で説明するように, FFT(高速フーリエ変換)を用いることによって計算量のオーダーを に改善することができます.
多項式補間を用いた解法
いま
相異なる個の点 を通る次多項式は一意に存在する
という多項式の性質を用い, 次のような手順でを計算することを考えます:
- Step 1: 相異なる個の点 を任意にとり, ()を計算する
- Step 2: ()を計算する
- Step 3: を通る一意な次多項式が求めるなので, 点の座標から を逆算する
ここでVandermonde行列 \begin{align} V = \left(\begin{array}{ccccc} 1 & x_0 & x_0^{2} & \ldots & x_0^{2n-2} \\ 1 & x_1 & x_1^{2} & \ldots & x_1^{2n-2} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 1 & x_{2n-2} & x_{2n-2}^{2} & \ldots & x_{2n-2}^{2n-2} \end{array}\right) \end{align} (次正方行列)を導入します. またとし, \begin{align} \mathbf{a} = \left(\begin{array}{c} a_0 \\ a_1 \\ \vdots \\ a_{2n-2} \end{array}\right), \ \mathbf{b} = \left(\begin{array}{c} b_0 \\ b_1 \\ \vdots \\ b_{2n-2} \end{array}\right) \end{align} とします.
するとstep 1の計算は \begin{align} \left(\begin{array}{c} f(x_0) \\ f(x_1) \\ \vdots \\ f(x_{2n-2}) \end{array}\right) = V \ \mathbf{a}, \hspace{1em} \left(\begin{array}{c} g(x_0) \\ g(x_1) \\ \vdots \\ g(x_{2n-2}) \end{array}\right) = V \ \mathbf{b} \end{align} のような行列の掛け算に帰着できます.
また Vandermonde行列の行列式が \begin{align} \det(V) = \prod_{0 \leq s < t\leq 2n - 2} (x_t - x_s) \end{align} であるという有名事実から, が相異なれば となり, は逆行列を持ちます. よってstep 3の計算も \begin{align} \left(\begin{array}{c} c_0 \\ c_1 \\ \vdots \\ c_{2n-2} \end{array}\right) = V^{-1} \ \left(\begin{array}{c} h(x_0) \\ h(x_1) \\ \vdots \\ h(x_{2n-2}) \end{array}\right) \end{align} という連立一次方程式の問題に帰着できます.
以上をまとめると, \begin{align} \left(\begin{array}{c} c_0 \\ c_1 \\ \vdots \\ c_{2n-2} \end{array}\right) = V^{-1} \ \left( V \ \mathbf{a} \circ V \ \mathbf{b} \right) \end{align} となります. ただし はベクトルの要素積です.
計算量
上の解法の計算量を各step毎に考えてみます.
前準備
まずstep 1, step 3の両方で使うVandermonde行列を先に求めておきます. 各に対し, を求めるには回の掛け算が必要になるので, 合計の計算回数は \begin{align} (2n - 1) \times (2n - 3) \end{align} 回程度になり, オーダーはになります
Step 1
それぞれに対して行列とベクトルの掛け算を行います. の各行に対してそれぞれ回スカラーの掛け算を行うので, 合計の計算回数は \begin{align} 2 \times (2n - 1) \times (2n - 1) \end{align} 回程度になり, オーダーはになります
Step 2
各 に対して を求めればよいので, 計算回数は回, オーダーはです
Step 3
変数の連立一次方程式をLU分解で解く場合の計算量は程度なので, である今回も計算量のオーダーはになります
以上をまとめると, 多項式補間を用いた解法の計算量は \begin{align} O(n^{2}) + O(n^{2}) + O(n) + O(n^{3}) = O(n^{3}) \end{align} 程度となり, 素直な解法よりも遅くなってしまいました*2.
しかし を任意にとってよいという性質を活かすことで, 各stepでの計算量を大幅にへらすことができます.
多項式補間 + DFT(離散フーリエ変換)を使った解法
以降 の定義域と値域をに拡張して考えます.
前準備の改良
まず標本点として ()をとってきます. ただし \begin{align} \omega_t = e^{2 \pi i / t} \end{align} です(は虚数単位). このとき \begin{align} \omega_{t}^{k} = \exp \left(\frac{2 \pi k}{t} i \right) = \cos \left(\frac{2 \pi k}{t} \right) + i \sin \left(\frac{2 \pi k}{t} \right) \end{align} となります.
たとえばのケースを複素平面上にプロットすると, 以下のようになります.
いまのべき乗は周期になっている, すなわち \begin{align} \forall \lambda \in \mathbb{Z}, \ \omega_t^{\lambda t + k} = \omega_t^{k} \end{align} が成り立つので, 上記のように定めた に対するVandermonde行列は, 次のように簡単に求めることができます:
とすると, \begin{align} V = \left(\begin{array}{ccccc} 1 & \omega^{0 \times 1} & \omega^{0 \times 2} & \ldots & \omega^{0 \times (2n - 2)} \\ 1 & \omega^{1 \times 1} & \omega^{1 \times 2} & \ldots & \omega^{1 \times (2n - 2)} \\ 1 & \omega^{2 \times 1} & \omega^{2 \times 2} & \ldots & \omega^{2 \times (2n - 2)} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 1 & \omega^{(2n - 2) \times 1} & \omega^{(2n - 2) \times 2} & \ldots & \omega^{(2n - 2) \times (2n - 2)} \end{array}\right) \\ = \left(\begin{array}{ccccc} 1 & \omega^{0} & \omega^{0} & \ldots & \omega^{0} \\ 1 & \omega^{1} & \omega^{2} & \ldots & \omega^{2n - 2} \\ 1 & \omega^{2} & \omega^{4} & \ldots & \omega^{2n - 3} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 1 & \omega^{2n - 2} & \omega^{2n - 3} & \ldots & \omega^{0} \end{array}\right) \end{align} すなわち の -要素()が \begin{align} V_{jk} = \omega^{k^{\prime}} \hspace{1em} \text{s.t.} \hspace{1em} k^{\prime} & \in \{ 0, 1, \ldots, 2n-2 \} \tag{1} \\ k^{\prime} & \equiv j \times k \text{ (mod $2n-1$)} \end{align} となります.
なおこのVandermonde行列 をDFT行列といい, をに左から掛けることを離散フーリエ変換, をに左から掛けることを離散フーリエ逆変換といいます(詳細は後述).
の要素に出てくる値は だけなので, 各に対して陽にその2乗, 3乗, ...を計算する必要はなくなりました. しかしmodの計算をする必要があるため, 計算量自体は のままです.
Step 3の改良
実は はユニタリ行列になっています. すなわち \begin{align} U^{*} U = U U^{*} = I \end{align} が成り立ちます(は単位行列で, は行列の共役転置を表します).
(証明)
まずの複素共役は, \begin{align} \bar{\omega^{k}} = \omega^{(2n - 1) - k} = \omega^{ - k} \end{align} となる(実数軸に関して線対称な位置への移動).
ここで行列の行目を切り出した行ベクトルを, 列目を切り出した列ベクトルをとする(). このとき \begin{align} U^{*}(j, \cdot) \ U(\cdot, k) & = \frac{1}{2n - 1} \left(\begin{array}{ccccc} (\bar{\omega^{j}})^{0} & (\bar{\omega^{j}})^{1} &(\bar{\omega^{j}})^{2} & \ldots & (\bar{\omega^{j}})^{2n - 2} \end{array}\right) \left(\begin{array}{c} \omega^{0 \times k} \\ \omega^{1 \times k} \\ \omega^{2 \times k} \\ \vdots \\ \omega^{(2n - 2) \times k} \end{array}\right) \\ & = \frac{1}{2n - 1} \left(\begin{array}{ccccc} \omega^{- j \times 0} & \omega^{- j \times 1} & \omega^{- j \times 2} & \ldots & \omega^{- j \times (2n - 2)} \end{array}\right) \left(\begin{array}{c} \omega^{0 \times k} \\ \omega^{1 \times k} \\ \omega^{2 \times k} \\ \vdots \\ \omega^{(2n - 2) \times k} \end{array}\right) \\ & = \frac{1}{2n - 1} \sum_{t=0}^{2n - 2} \omega^{t \times (k - j)} \\ & = \begin{cases} \cfrac{1}{2n - 1} \cfrac{1 - (\omega^{k - j})^{2n - 1} }{1 - \omega^{k - j}} & \text{if $j \neq k$}\\ \cfrac{1}{2n - 1} (\sum_{t=0}^{2n - 2} 1) & \text{if $j = k$} \end{cases} \\ & = \begin{cases} 0 & \text{if $j \neq k$}\\ 1 & \text{if $j = k$} \end{cases} \end{align} のケースの最後の変形は, (mod )より, \begin{align} (\omega^{k - j})^{2n - 1} = \omega^{0} = 1 \end{align} となることを用いた.
に関しても, 全く同様にして示すことができる.
したがって \begin{align} V^{-1} = \frac{1}{\sqrt{2n - 1}} U^{-1} = \frac{1}{\sqrt{2n - 1}} U^{ * } = \frac{1}{2n - 1} V^{ * } \end{align} となり, の逆行列を計算する処理は, の共役転置を取る操作に置き換えられます. これを用いるとstep 3の計算量は
- の計算に
- とベクトル の掛け算に
かかるので, 合わせて に改善します.
したがって前処理〜step 3までの全体の計算量のオーダーは \begin{align} O(n^{2}) + O(n^{2}) + O(n) + O(n^{2}) = O(n^{2}) \end{align} となります. ただこれでもまだ「単純な方法」と同じオーダーになってしまう(かつこちらの方が複雑な処理をする分時間がかかる)ので, さらに改良してのオーダーを目指します.
余談: DFTについて
離散フーリエ変換の意味について少し補足をしておきます. とし, のDFT行列をとします. を離散フーリエ変換した結果がであるとすると, 上の議論から \begin{align} \mathbf{y} = F_k \mathbf{x} & \Longleftrightarrow F_k^{- 1} \mathbf{y} = I \mathbf{x} \\ & \Longleftrightarrow \cfrac{1}{k} F_k^{*} \mathbf{y} = I \mathbf{x} \\ & \Longleftrightarrow \cfrac{1}{k} \ \left(\begin{array}{ccccc} 1 & 1 & 1 & \ldots & 1 \\ \omega^{- 0} & \omega^{- 1} & \omega^{- 2} & \ldots & \omega^{- (k - 1)} \\ \omega^{- 0} & \omega^{- 2} & \omega^{- 4} & \ldots & \omega^{- (k - 1) \times 2} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ \omega^{- 0} & \omega^{- (k - 1)} & \omega^{- (k - 1) \times 2} & \ldots & \omega^{- (k - 1) \times (k - 1)} \end{array}\right) \mathbf{y} = I \mathbf{x} \\ & \Longleftrightarrow y_0 \phi_0 + \ldots + y_{k - 1} \phi_{k - 1} = x_0 e_0 + \ldots x_{k - 1} e_{k - 1} \end{align} が成り立ちます. ただし \begin{align} \phi_j = \cfrac{1}{k} \ \left(\begin{array}{c} \omega^{- j \times 0} \\ \omega^{- j \times 1} \\ \omega^{- j \times 2} \\ \vdots \\ \omega^{- j \times (k - 1)} \end{array}\right) \end{align} であり, は番目の要素が1, それ以外が0のベクトルです.
したがって離散フーリエ変換は, 標準基底での座標を, を基底としたときの座標に, 座標変換する操作になります.
なお であることから, はの正規直交基底になります.
多項式補間 + FFTを使った解法
分割統治法を使って離散フーリエ変換を高速化しようというのが, FFT(高速フーリエ変換)の考え方です.
いま は2のべき乗 であると仮定し, をフーリエ変換した結果である を求めることを考えます(が2のべき乗でない場合は, よりも大きい2のべき乗の中で一番小さいものをと再定義し, それに合わせてに0を付け足せばよいです).
まず事実として, 2の倍数であるようなに対しては \begin{align} \omega_t^{2} = \exp \left( 2 \times \frac{2 \pi}{t} i \right) = \exp \left( \frac{2 \pi}{t / 2} i \right) = \omega_{t / 2} \tag{2} \end{align} が成り立ちます.
次に行列の偶数番号列だけを抜き出した行列を, 奇数番号列だけを抜き出した行列を とします. すなわち \begin{align} F_k^{\text{even}} & = \left(\begin{array}{ccccc} 1 & \omega^{0 \times 2} & \omega^{0 \times 4} & \ldots & \omega^{0 \times (k - 2)} \\ 1 & \omega^{1 \times 2} & \omega^{1 \times 4} & \ldots & \omega^{1 \times (k - 2)} \\ 1 & \omega^{2 \times 2} & \omega^{2 \times 4} & \ldots & \omega^{2 \times (k - 2)} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 1 & \omega^{(k - 1) \times 2} & \omega^{(k - 1) \times 4} & \ldots & \omega^{(k - 1) \times (k - 2)} \end{array}\right) \\ F_k^{\text{odd}} & = \left(\begin{array}{cccc} \omega^{0 \times 1} & \omega^{0 \times 3} & \ldots & \omega^{0 \times (k - 1)} \\ \omega^{1 \times 1} & \omega^{1 \times 3} & \ldots & \omega^{1 \times (k - 1)} \\ \omega^{2 \times 1} & \omega^{2 \times 3} & \ldots & \omega^{2 \times (k - 1)} \\ \vdots & \vdots & \ddots & \vdots \\ \omega^{(k - 1) \times 1} & \omega^{(k - 1) \times 3} & \ldots & \omega^{(k - 1) \times (k - 1)} \end{array}\right) \end{align} とし, ] と定義します.
同様にの偶数番号行だけ抜き出したベクトルを, 奇数番号行だけ抜き出したベクトルをそれぞれとし, \begin{align} \mathbf{x}^{\prime} = \left(\begin{array}{c} \mathbf{x}^{\text{even}} \\ \mathbf{x}^{\text{odd}} \end{array}\right) \end{align} とします.
このとき定義より \begin{align} \mathbf{y} = F_k^{\prime} \mathbf{x}^{\prime} \end{align} が成り立ちます. からを計算するコストはなので, の演算さえで実現できれば, 全体の計算量もにできます.
いまは \begin{align} F_k^{\text{even}} = \left(\begin{array}{ccc} 1 & \ldots & \omega_k^{0 \times (k - 2)} \\ \vdots & \ddots & \vdots \\ 1 & \ldots & \omega_k^{(k/2 - 1) \times (k - 2)} \\ \hline 1 & \ldots & \omega_k^{k/2 \times (k - 2)} \\ \vdots & \ddots & \vdots \\ 1 & \ldots & \omega_k^{(k - 1) \times (k - 2)} \end{array}\right) & = \left(\begin{array}{ccc} 1 & \ldots & \omega_{k/2}^{0 \times (k/2 - 1)} \\ \vdots & \ddots & \vdots \\ 1 & \ldots & \omega_{k/2}^{(k/2 - 1) \times (k/2 - 1)} \\ \hline 1 & \ldots & \omega_{k/2}^{k/2 \times (k/2 - 1)} \\ \vdots & \ddots & \vdots \\ 1 & \ldots & \omega_{k/2}^{(k - 1) \times (k/2 - 1)} \end{array}\right) \hspace{0.5em} \text{($\because$ (2))} \\ & = \left(\begin{array}{ccc} 1 & \ldots & \omega_{k/2}^{0 \times (k/2 - 1)} \\ \vdots & \ddots & \vdots \\ 1 & \ldots & \omega_{k/2}^{(k/2 - 1) \times (k/2 - 1)} \\ \hline 1 & \ldots & \omega_{k/2}^{0 \times (k/2 - 1)} \\ \vdots & \ddots & \vdots \\ 1 & \ldots & \omega_{k/2}^{(k/2 - 1) \times (k/2 - 1)} \end{array}\right) \hspace{0.5em} \text{($\because$ (1))} \\ & = \left(\begin{array}{c} F^{\prime}_{k/2} \\ F^{\prime}_{k/2} \end{array}\right) \end{align} と変形できます. 一方とすると, も \begin{align} F_k^{\text{odd}} & = \left(\begin{array}{ccc} \omega_k^{0 \times 1} & \ldots & \omega_k^{0 \times (k - 1)} \\ \vdots & \ddots & \vdots \\ \omega_k^{(k/2 - 1) \times 1} & \ldots & \omega_k^{(k/2 - 1) \times (k - 1)} \\ \hline \omega_k^{k/2 \times 1} & \ldots & \omega_k^{k/2 \times (k - 1)} \\ \vdots & \ddots & \vdots \\ \omega_k^{(k - 1) \times 1} & \ldots & \omega_k^{(k - 1) \times (k - 1)} \end{array}\right) \\ & = \left(\begin{array}{ccc} \omega_k^{0 \times 1} & \ldots & \omega_k^{0 \times (k - 1)} \\ \vdots & \ddots & \vdots \\ \omega_k^{(k/2 - 1) \times 1} & \ldots & \omega_k^{(k/2 - 1) \times (k - 1)} \\ \hline \omega_k^{(k - k/2) \times 1} & \ldots & \omega_k^{(k - k/2) \times (k - 1)} \\ \vdots & \ddots & \vdots \\ \omega_k^{(k - (k - 1) ) \times 1} & \ldots & \omega_k^{(k - (k - 1) ) \times (k - 1)} \end{array}\right) \\ & = \left(\begin{array}{ccc} \omega_k^{0 \times 1} & \ldots & \omega_k^{0 \times (k - 1)} \\ \vdots & \ddots & \vdots \\ \omega_k^{(k/2 - 1) \times 1} & \ldots & \omega_k^{(k/2 - 1) \times (k - 1)} \\ \hline - \omega_k^{0 \times 1} & \ldots & - \omega_k^{0 \times (k - 1)} \\ \vdots & \ddots & \vdots \\ - \omega_k^{(k/2 - 1) \times 1} & \ldots & - \omega_k^{(k/2 - 1) \times (k - 1)} \end{array}\right) \\ & = \left(\begin{array}{c} D F^{\prime}_{k/2} \\ \hline - D F^{\prime}_{k/2} \end{array}\right) \end{align} と変形できます. 3番目の等式は, 複素単位円上で原点対称な位置への移動を行っていることから来ています.
まとめると \begin{align} F_k = [F_k^{\text{even}}, F_k^{\text{odd}}] = \left(\begin{array}{cc} F_{k/2} & D F_{k/2} \\ F_{k/2} & - D F_{k/2} \end{array}\right) \end{align} となるので, \begin{align} \left(\begin{array}{c} y_{0} \\ \vdots \\ y_{k/2 - 1} \end{array} \right) & = F_{k/2} \mathbf{x}^{\text{even}} + D F_{k/2} \mathbf{x}^{\text{odd}} \\ \left(\begin{array}{c} y_{k/2} \\ \vdots \\ y_{k - 1} \end{array} \right) & = F_{k/2} \mathbf{x}^{\text{even}} - D F_{k/2} \mathbf{x}^{\text{odd}} \end{align} が成り立ちます. よって次元ベクトルにFFTを行う操作は,
- を奇数列と偶数列に分割する操作()
- 次元ベクトルに対しFFTを行う計算2つ(計算量は後述)
- の計算(が対角行列なのでで可能)
- ベクトルの和/差を計算する操作()
に分割できます
FFTの計算量
が次元の時にFFTを行う計算量をとします. すると上の分割計算の手順より, 十分大きなに対して \begin{align} T(k) \leq 2 T(k/2) + c k \end{align} という関係が成り立ちます(). また1点に対するFFTは元の値そのものなので, です. この時, \begin{align} T(k) = O(k \log k) \end{align} が成り立ちます(帰納法で証明できます).
高速フーリエ逆変換に関しても同様の分割統治法で計算でき, 計算量のオーダーはです.
アルゴリズム全体の計算量
長々と説明してきましたが, 結局のところを求める手順としては
となります. したがってトータルの計算量はとなり, 単純な方法の場合よりも少ない計算量を実現することができました.
実装
これをPythonで実装すると, 以下のようになります.
def get_poly_product_fft(f_coefs, g_coefs): """ Parameters: f_coefs: array_like. fの係数のリスト. [0次の係数, 1次の係数, ..., 最高次の係数] g_coefs: array_like. gの係数のリスト. [0次の係数, 1次の係数, ..., 最高次の係数] Returns: h_coefs: ndarray. h=fgの係数のリスト. [0次の係数, 1次の係数, ..., 最高次の係数] """ h_size = len(f_coefs) + len(g_coefs) - 1 h_fft_coefs = np.fft.fft(f_coefs, n=h_size) * np.fft.fft(g_coefs, n=h_size) h_coefs = np.fft.ifft(h_fft_coefs) return np.real(h_coefs)
例題を解いてみると,
get_poly_product_fft([3, 2, 1], [4, 3]) >> array([12., 17., 10., 3.])
となり, 手で解いた結果と一致しました.
速度比較: 単純な方法 vs 多項式補間 + FFTを使った解法
私は「多項式補間 + FFT」でを求める方法を初めて聞いたとき, 「いくら計算量のオーダーが小さくなるといっても, こんなに複雑なことをして計算が速くなるのか?」と疑問に思いました(はじめて聞いた人は普通そう思うでしょう).
結論: FFTしゅごい.