AtCoderのFFT問題に挑戦

この記事は「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;
}