この記事は「IQ1の2まいめっ Advent Calendar 2018」の9日目の記事です。
AtCoderというプログラミングコンテストサイトに、「高速フーリエ変換」という、タイトルどおり高速フーリエ変換 (fast Fourier transform, FFT) を題材とした問題があります。この記事では、その問題を通じて、高速フーリエ変換について説明していきます。
問題と高速フーリエ変換の関係
扱う問題は、簡単に言い換えると「配列 $A_0, A_1, \ldots , A_N ; (A_0 = 0)$ と $B_0, B_1, \ldots , B_N ; (B_0 = 0)$ があるとき、各 $k ; (1 \le k \le 2N)$ に対して以下の $C_k$ を求めよ」となります。
C_k = \sum_{i=0}^k A_i B_{k-i}
ここで、$C_k$ は多項式 $p_A(x) = A_0 + A_1 x + \ldots + A_N x^N, ; p_B(x) = B_0 + B_1 x + \ldots + B_N x^N$ の積 $p_A(x) p_B(x)$ の $k$ 次の項の係数と同値です。よって、この問題は$p_A(x)$ と $p_B(x)$ の積 $p_A(x) p_B(x)$ を求めることができればよいです。
$p_A(x) p_B(x)$ を多項式の筆算の要領で求める場合、$A_i B_i$ を各$i, ; j$ に対して計算し、それらをもとに各 $k$ に対して $C_k$ を計算する必要があるため、計算量は $O (N^2)$ です。しかし、本記事で扱っているAtCoderの「高速フーリエ変換」という問題の制限時間は5secであり、より計算量の少ない方法を用いる必要があります。そこで、$p_A(x) p_B(x)$ をより効率よく計算するために高速フーリエ変換 (fast Fourier transform, FFT) と呼ばれる手法を用います。
高速な多項式乗算への準備
高速フーリエ変換により多項式乗算を高速に計算することができますが、そのことを理解するためにはまず、多項式補間と離散フーリエ変換、高速フーリエ変換について説明する必要があります。
多項式補間
$N$ 次多項式 $p(x) = A_0 + A_1 x + \ldots + A_N x^N ; (A_N \ne 0)$ に対し、 $y_i = p(x_i) ; (0 \le i \le N)$ となるような $x_i, y_i; (0 \le i \le N)$ を考えます。ただし、各 $x_i$ は異なる値とします。
このとき、
\begin{bmatrix}
y_0 \\
y_1 \\
\vdots \\
y_N
\end{bmatrix}
=
\begin{bmatrix}
A_0 + A_1 x_0 + \ldots + A_N x_0^N \\
A_0 + A_1 x_1 + \ldots + A_N x_1^N \\
\vdots \\
A_0 + A_1 x_N + \ldots + A_N x_N^N
\end{bmatrix}
=
\begin{bmatrix}
1 & x_0 & \ldots & x_0^N \\
1 & x_1 & \ldots & x_1^N \\
\vdots & \vdots & \ddots & \vdots \\
1 & x_N & \ldots & x_N^N
\end{bmatrix}
\begin{bmatrix}
A_0 \\
A_1 \\
\vdots \\
A_N
\end{bmatrix}
となります。ここで、
\boldsymbol{y} =
\begin{bmatrix}
y_0 \\
y_1 \\
\vdots \\
y_N
\end{bmatrix}
,
X =
\begin{bmatrix}
1 & x_0 & \ldots & x_0^N \\
1 & x_1 & \ldots & x_1^N \\
\vdots & \vdots & \ddots & \vdots \\
1 & x_N & \ldots & x_N^N
\end{bmatrix}
,
\boldsymbol{A} =
\begin{bmatrix}
A_0 \\
A_1 \\
\vdots \\
A_N
\end{bmatrix}
とおくと、
\mathrm{det} X = \prod_{0 \le i < j \le n} (x_j - x_i)
が成り立つことが知られています。各 $x_i$ が異なる値をとると仮定しているので、$\mathrm{det} X \ne 0$ であり、$X$ には逆行列が存在します。よって、
\boldsymbol{A} = X^{-1} \boldsymbol{y}
が成り立ちます。このことより、$N$ 次多項式 $p(x) = A_0 + A_1 x + \ldots + A_N x^N$ の係数は $N + 1$ 個の異なる点 $x_0, x_1, \ldots, x_N$ とそれらに対応する $p(x)$ の値 $p(x_0), p(x_1), \ldots, p(x_N)$ を用いて表現できることがわかります。このことは、$N$ 次多項式 $p(x)$ は $N + 1$ 個の異なる点 $x_0, x_1, \ldots, x_N$ とそれらに対応する $p(x)$ の値 $p(x_0), p(x_1), \ldots, p(x_N)$ により一意に定まるということを表していて、これにより $p(x)$ を求めることを多項式補間といいます。
離散フーリエ変換と高速フーリエ変換
離散フーリエ変換 (discrete Fourier transform, DFT) とは、関数 $f(x)$ を以下の関数 $F(t)$ に変換することをいいます。
F(t) = \sum_{x=0}^{n-1} f(x) \exp \left( -i \frac{2 \pi t x}{n} \right)
また、逆に $F(t)$ を $f(x)$ に変換する操作を逆離散フーリエ変換 (inverse discrete Fourier transform, IDFT) といい、以下のように表されます。
f(x) = \frac{1}{n} \sum_{t=0}^{n-1} F(t) \exp \left( i\frac{2 \pi x t}{n} \right)
これら、離散フーリエ変換(または逆離散フーリエ変換)を高速に計算する高速フーリエ変換(または逆離散フーリエ変換)ですが、以降の説明では $n$ が2の累乗であると仮定します(本当は $n$ が2の累乗でなくても高速フーリエ変換はできるそうなのですが、今回は簡単のために $n$ が2の累乗であるとします)。$n$ が2の累乗でない場合は、$n \le 2^m$ となるような $2^m$ まで拡張し、$f(x) = 0 ; (n \le x \le 2^m - 1)$ とすれば以下に説明する方法で高速フーリエ変換をすることができます。
各 $0 \le t \le n - 1$ について、$F(t)$ を求めることを考えます。ここで、$w_n = \exp( -i 2 \pi / n)$ とおくと以下の式が成り立ちます。
\begin{bmatrix}
F(0) \\
F(1) \\
\vdots \\
F(n - 1)
\end{bmatrix}
=
\begin{bmatrix}
f(0) w_n^0 + f(1) w_n^0 + \ldots + f(n - 1) w_n^0 \\
f(0) w_n^0 + f(1) w_n^1 + \ldots + f(n - 1) w_n^{n-1} \\
\vdots \\
f(0) w_n^0 + f(1) w_n^{(n-1) 1} + \ldots + f(n - 1) w_n^{(n-1) (n-1)}
\end{bmatrix}
\\ =
\begin{bmatrix}
w_n^0 & w_n^0 & \ldots & w_n^0 \\
w_n^0 & w_n^1 & \ldots & w_n^{n-1} \\
\vdots & \vdots & \ddots & \vdots \\
w_n^0 & w_n^{n-1} & \ldots & w_n^1
\end{bmatrix}
\begin{bmatrix}
f(0) \\
f(1) \\
\vdots \\
f(n - 1)
\end{bmatrix}
ここで、
\boldsymbol{F} =
\begin{bmatrix}
F(0) \\
F(1) \\
\vdots \\
F(n - 1)
\end{bmatrix}
, \;
\boldsymbol{f} =
\begin{bmatrix}
f(0) \\
f(1) \\
\vdots \\
f(n - 1)
\end{bmatrix}
, \;
W_n =
\begin{bmatrix}
w_n^0 & w_n^0 & \ldots & w_n^0 \\
w_n^0 & w_n^1 & \ldots & w_n^{n-1} \\
\vdots & \vdots & \ddots & \vdots \\
w_n^0 & w_n^{n-1} & \ldots & w_n^1
\end{bmatrix}
とおくと、先ほどの式は
\boldsymbol{F} = W_n \: \boldsymbol{f}
となります。この式は以下のように書き換えられます。
\begin{bmatrix}
F(0) \\
F(1) \\
\vdots \\
F(n / 2 - 1) \\
F(n / 2) \\
F(n / 2 + 1) \\
\vdots \\
F(n - 1)
\end{bmatrix}
=
\begin{bmatrix}
f(0) w_n^0 + f(1) w_n^0 + \ldots + f(n - 1) w_n^0 \\
f(0) w_n^0 + f(1) w_n^1 + \ldots + f(n - 1) w_n^{n-1} \\
\vdots \\
f(0) w_n^0 + f(1) w_n^{(n/2-1) 1} + \ldots + f(n - 1) w_n^{(n/2-1) (n-1)} \\
f(0) w_n^0 + f(1) w_n^{(n/2) 1} + \ldots + f(n - 1) w_n^{(n/2) (n-1)} \\
f(0) w_n^0 + f(1) w_n^{(n/2 + 1) 1} + \ldots + f(n - 1) w_n^{(n/2+1)(n-1)} \\
\vdots \\
f(0) w_n^0 + f(1) w_n^{(n-1) 1} + \ldots + f(n - 1) w_n^{(n-1) (n-1)}
\end{bmatrix}
\\ =
\begin{bmatrix}
w_n^0 & w_n^0 & \ldots & w_n^0 & w_n^0 & w_n^0 & \ldots & w_n^0 \\
w_n^0 & w_n^2 & \ldots & w_n^{n-2} & w_n^1 & w_n^3 & \ldots & w_n^{n-1} \\
\vdots & \vdots & \ddots & \vdots & \vdots & \vdots & \ddots & \vdots \\
w_n^0 & w_n^{n-2} & \ldots & w_n^2 & w_n^{n/2-1} & w_n^{(n/2-1) + (n-2)} & \ldots & w_n^{(n/2 - 1) + 2} \\
w_n^0 & w_n^0 & \ldots & w_n^0 & w_n^{n/2} & w_n^{n/2} & \ldots & w_n^{n/2}\\
w_n^0 & w_n^2 & \ldots & w_n^{n-2} & w_n^{n/2 + 1} & w_n^{n/2+3} & \ldots & w_n^{n/2 + n - 1} \\
\vdots & \vdots & \ddots & \vdots & \vdots & \vdots & \ddots & \vdots \\
w_n^0 & w_n^{n-2} & \ldots & w_n^2 & w_n^{n/2 + (n/2-1)} & w_n^{n/2 + (n/2-1) + n -2} & \ldots & w_n^{n/2+(n/2-1)+2}
\end{bmatrix}
\begin{bmatrix}
f(0) \\
f(2) \\
\vdots \\
f(n - 2) \\
f(1) \\
f(3) \\
\vdots \\
f(n - 1)
\end{bmatrix}
ここで、
\boldsymbol{F}_0 =
\begin{bmatrix}
F(0) \\
F(1) \\
\vdots \\
F(n/2 - 1)
\end{bmatrix}
, \;
\boldsymbol{F}_1 =
\begin{bmatrix}
F(n/2) \\
F(n/2 + 1) \\
\vdots \\
F(n - 1)
\end{bmatrix}
, \;
\boldsymbol{f}_{\mathrm{even}} =
\begin{bmatrix}
f(0) \\
f(2) \\
\vdots \\
f(n - 2)
\end{bmatrix}
, \;
\boldsymbol{f}_{\mathrm{odd}} =
\begin{bmatrix}
f(1) \\
f(3) \\
\vdots \\
f(n - 1)
\end{bmatrix}
, \;
D_n = \mathrm{diag}\left( w_{2n}^0, w_{2n}^1, \ldots, w_{2n}^{n-1} \right)
とおくと、先ほどの式は以下のように表せます。
\begin{bmatrix}
\boldsymbol{F}_0 \\
\boldsymbol{F}_1
\end{bmatrix}
=
\begin{bmatrix}
W_{n/2} & D_{n/2} W_{n/2} \\
W_{n/2} & w_n^{n/2} D_{n/2} W_{n/2}
\end{bmatrix}
\begin{bmatrix}
\boldsymbol{f}_{\mathrm{even}} \\
\boldsymbol{f}_{\mathrm{odd}}
\end{bmatrix}
\\ \therefore
\boldsymbol{F} =
W_n \: \boldsymbol{f}
=
\begin{bmatrix}
W_{n/2} \: \boldsymbol{f}_{\mathrm{even}} + D_{n/2} W_{n/2} \: \boldsymbol{f}_{\mathrm{odd}} \\
W_{n/2} \: \boldsymbol{f}_{\mathrm{even}} - D_{n/2} W_{n/2} \: \boldsymbol{f}_{\mathrm{odd}} \\
\end{bmatrix}
この式から、$\boldsymbol{F}$ を求めるには分割統治法を用いることができるとわかります。このようにして、$t = 0, 1, \ldots, n - 1$ の $n$ 点における関数 $f(x)$ の離散フーリエ変換 $F(t)$ の値を、計算量 $O(n \log n)$ で求めることができます。これを高速フーリエ変換といいます。同様にして、逆離散フーリエ変換を行うことを逆高速フーリエ変換 (inverse fast Fourier transform, IFFT) といいます。
高速フーリエ変換を用いた多項式乗算
$n - 1$ 次多項式 $p_A(x) = A_0 + A_1 x + \ldots + A_{n-1} x^{n-1}$ と、$f(0) = A_0, f(1) = A_1, \ldots, f(n-1) = A_{n-1}$ となるような関数 $f(x)$ を考えます。すると、$t = 0, 1, \ldots, n-1$ の $n$ 点における関数 $f(x)$ の離散フーリエ変換 $F(t)$ の値を求める式は、行列を用いた表現で以下のように表せます。
\begin{bmatrix}
F(0) \\
F(1) \\
\vdots \\
F(n -1)
\end{bmatrix}
=
\begin{bmatrix}
w_{n}^0 & w_{n}^0 & \ldots & w_{n}^0 \\
w_{n}^0 & w_{n}^1 & \ldots & w_{n}^{n-1} \\
\vdots & \vdots & \ddots & \vdots \\
w_{n}^{0} & w_{n}^{n-1} & \ldots & w_{n}^{(n-1) (n-1)}
\end{bmatrix}
\begin{bmatrix}
A_0 \\
A_1 \\
\vdots \\
A_{n-1}
\end{bmatrix}
$F(t) = A_0 + A_1 w_{n}^t + \ldots + A_{n-1} \left( w_{n}^t \right)^{n-1}$ より、$F(t)$ は $p_A \left( w_{n}^t \right)$ に他なりません。高速フーリエ変換により、$x = w_{n}^0, w_{n}^1, \ldots, w_{n}^{n-1}$ の $n$ 点における $n$ 次多項式 $p_A(x)$ の値は $O(n \log n)$ で求めることができます。同様にして、$x = w_{n}^0, w_{n}^1, \ldots, w_{n}^{n-1}$ の $n$ 点における $n - 1$ 次多項式 $p_B(x) = B_0 + B_1 x + \ldots + B_{n-1} x^{n-1}$ の値も $O(n \log n)$ で求めることができます。
ここで、$p_C(x) = p_A(x) p_B(x) = C_0 + C_1 x + \ldots + C_{2 n - 2} x^{2 n - 2}$ の係数 $C_k ; (0 \le k \le 2n - 2)$ を求めることを考えます。$p_A(x)$ の係数を拡張し、$A_k = 0 ; (n \le k \le 2n-1)$ とした状態で以下のように高速フーリエ変換をしておきます。$p_B(x)$ も同様にして $F_B(0), F_B(1), \ldots, F_B(2n - 1)$ を求めておきます。
\begin{bmatrix}
F_A(0) \\
F_A(1) \\
\vdots \\
F_A(2n -1)
\end{bmatrix}
=
\begin{bmatrix}
w_{2n}^0 & w_{2n}^0 & \ldots & w_{2n}^0 \\
w_{2n}^0 & w_{2n}^1 & \ldots & w_{2n}^{2n-1} \\
\vdots & \vdots & \ddots & \vdots \\
w_{2n}^{0} & w_{2n}^{2n-1} & \ldots & w_{2n}^{(2n-1) (2n-1)}
\end{bmatrix}
\begin{bmatrix}
A_0 \\
A_1 \\
\vdots \\
A_{2n-1}
\end{bmatrix}
$p_C\left(w_{2n}^t \right) = p_B\left(w_{2n}^t \right) p_B\left(w_{2n}^t \right) = F_A(t) F_B(t)$ より、以下の式が成り立ちます(なお、$p_C(x)$ の係数ではない $C_{2n-1}$ も便宜的に含めています)。
\begin{bmatrix}
F_A(0) F_B(0) \\
F_A(1) F_B(1) \\
\vdots \\
F_A(2n -1) F_B(2n-1)
\end{bmatrix}
=
\begin{bmatrix}
w_{2n}^0 & w_{2n}^0 & \ldots & w_{2n}^0 \\
w_{2n}^0 & w_{2n}^1 & \ldots & w_{2n}^{2n-1} \\
\vdots & \vdots & \ddots & \vdots \\
w_{2n}^{0} & w_{2n}^{2n-1} & \ldots & w_{2n}^{(2n-1) (2n-1)}
\end{bmatrix}
\begin{bmatrix}
C_0 \\
C_1 \\
\vdots \\
C_{2n-1}
\end{bmatrix}
\\ \therefore
\begin{bmatrix}
C_0 \\
C_1 \\
\vdots \\
C_{2n-1}
\end{bmatrix}
=
\begin{bmatrix}
w_{2n}^0 & w_{2n}^0 & \ldots & w_{2n}^0 \\
w_{2n}^0 & w_{2n}^1 & \ldots & w_{2n}^{2n-1} \\
\vdots & \vdots & \ddots & \vdots \\
w_{2n}^{0} & w_{2n}^{2n-1} & \ldots & w_{2n}^{(2n-1) (2n-1)}
\end{bmatrix}^{-1}
\begin{bmatrix}
F_A(0) F_B(0) \\
F_A(1) F_B(1) \\
\vdots \\
F_A(2n -1) F_B(2n-1)
\end{bmatrix}
\\ = \frac{1}{2n}
\begin{bmatrix}
w_{2n}^0 & w_{2n}^0 & \ldots & w_{2n}^0 \\
w_{2n}^0 & w_{2n}^{-1} & \ldots & w_{2n}^{-(2n-1)} \\
\vdots & \vdots & \ddots & \vdots \\
w_{2n}^{0} & w_{2n}^{-(2n-1)} & \ldots & w_{2n}^{-1}
\end{bmatrix}
\begin{bmatrix}
F_A(0) F_B(0) \\
F_A(1) F_B(1) \\
\vdots \\
F_A(2n -1) F_B(2n-1)
\end{bmatrix}
$C_0, C_1, \ldots, C_{2n-2} ; (, C_{2n-1})$ は逆高速フーリエ変換により $O(n \log n)$ で求められます。
以上から、$n-1$ 次多項式 $p_A(x) = A_0 + A_1 x + \ldots + A_{n-1} x^{n-1}, ; p_B(x) = B_0 + B_1 x + \ldots + B_{n-1} x^{n-1}$ の多項式乗算は $O(n \log n)$ で計算できます。
よって、AtCoderの「高速フーリエ変換」問題は $O(N \log N)$ で解くことができます。
C言語で実装
参考までに、C言語で書いたAtCoderでACだったコードを以下に示します。
#include <stdio.h>
#include <complex.h>
#include <math.h>
void _fft(double complex a[], int n, int sign) {
if (n == 1)
return;
int n_half = n >> 1;
double complex a0[n_half], a1[n_half];
for (int i = 0; i < n_half; i++) {
a0[i] = a[i * 2];
a1[i] = a[i * 2 + 1];
}
_fft(a0, n_half, sign);
_fft(a1, n_half, sign);
double complex w = 1;
double complex w0 = cexp(M_PI * sign * I / n_half);
for (int i = 0; i < n_half; i++) {
double complex t = w * a1[i];
a[i] = a0[i] + t;
a[n_half + i] = a0[i] - t;
w *= w0;
}
}
void fft(double complex a[], int n) {
_fft(a, n, -1);
}
void inverse_fft(double complex a[], int n) {
_fft(a, n, 1);
for (int i = 0; i < n; i++) {
a[i] /= n;
}
}
int main() {
int n;
scanf("%d", &n);
int n2 = 1;
for (; n2 <= (n << 1); n2 <<= 1) {}
double complex a[n2], b[n2];
a[0] = b[0] = 0;
for (int i = 1; i <= n; i++) {
double x, y;
scanf("%lf %lf", &x, &y);
a[i] = x;
b[i] = y;
}
for (int i = n + 1; i < n2; i++) { // zero padding
a[i] = b[i] = 0;
}
fft(a, n2);
fft(b, n2);
for (int i = 0; i < n2; i++) {
a[i] *= b[i];
}
inverse_fft(a, n2);
for (int i = 1; i <= (n << 1); i++) {
printf("%d\n", (int) round(creal(a[i])));
}
return 0;
}
参考にしたページ
- https://atc001.contest.atcoder.jp/tasks/fft_c
- https://www.geeksforgeeks.org/fast-fourier-transformation-poynomial-multiplication/
- http://hooktail.org/computer/index.php?%B9%E2%C2%AE%A5%D5%A1%BC%A5%EA%A5%A8%CA%D1%B4%B9
- https://ja.wikipedia.org/wiki/%E5%A4%9A%E9%A0%85%E5%BC%8F%E8%A3%9C%E9%96%93
- https://ja.wikipedia.org/wiki/%E3%83%B4%E3%82%A1%E3%83%B3%E3%83%87%E3%83%AB%E3%83%A2%E3%83%B3%E3%83%89%E3%81%AE%E8%A1%8C%E5%88%97%E5%BC%8F