Go to Qiita Advent Calendar 2024 Top
LoginSignup
7
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

More than 5 years have passed since last update.

@hi-watana

AtCoderのFFT問題に挑戦

Last updated at Posted at 2018-12-08

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

参考にしたページ

7
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
7
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?