離散フーリエ変換と数論変換 (3) 高速フーリエ変換 (FFT)

この連載について

離散フーリエ変換 (DFT) および数論変換 (NTT) の原理、そしてそれらのプログラミングにおける実装方法について記述します。
各用語の定義の違いを明確にするため、それぞれについての回を分割します。

1. 離散フーリエ変換 (DFT)
2. 数列の畳み込み
3. 高速フーリエ変換 (FFT) (本稿)
4. ビット反転置換
5. 数論変換 (NTT)
6. NTT の高速化
7. 中国剰余定理による値の復元

本文

今回は、離散フーリエ変換 (DFT) において演算を効率化する高速フーリエ変換 (fast Fourier transform, FFT) と、そのプログラミングについて記述します。

前回に引き続き、多項式をベクトルとして考えることにします。
すなわち、$n-1$次多項式は$\mathbf{C}^{~n}$に属します。

DFT の定義を再確認しておきます。

実験

ここでは逆変換である IDFT については考えず、DFT の高速化を考えます。
すなわち、与えられた$(c_k)$に対して$(f(\omega_n^{~k}))$を効率的に求める方法を考察していきます。
まず、$n=4$として実験してみましょう。

$$ ~~ f(x) = c_0 + c_1 x + c_2 x^2 + c_3 x^3 $$
として、各$f(\omega_4^{~k})$を変形してみると、
$$ ~~ \mathcal{F}_4(f) = \begin{pmatrix} f(1) \\ f(i) \\ f(-1) \\ f(-i) \end{pmatrix} = \begin{pmatrix} (c_0+c_2) + 1 \cdot (c_1+c_3) \\ (c_0-c_2) + i \cdot (c_1-c_3) \\ (c_0+c_2) - 1 \cdot (c_1+c_3) \\ (c_0-c_2) - i \cdot (c_1-c_3) \end{pmatrix} $$
この式を観察すると、例えば$c_0+c_2$や$i \cdot (c_1-c_3)$など共通する部分を再利用できそうです。またこの部分が、長さが半分の DFT である$\mathcal{F}_2$の形となっていることから、効果的な短縮ができそうだと予想できます。

実際、
$$ ~~ g_0(x) := c_0 + c_2 x \\ ~~ g_1(x) := c_1 + c_3 x $$
とすると、
$$ ~~ f(x) = c_0 + c_2 x^2 + x(c_1 + c_3 x^2) = g_0(x^2) + x g_1(x^2) $$
の形式となります。
また、$(\omega_4^{~k})^2 = \omega_2^{~k}$であることから、$\mathcal{F}_2(g_0)$および$\mathcal{F}_2(g_1)$の各成分が現れます。

高速フーリエ変換

前節における考察を一般化します。
$n$を$2$の冪とします。

$n>1$のとき、
$$ ~~ g_0(x) := c_0 + c_2 x + c_4 x^2 + \cdots = \sum_{0 \leq k < n/2} c_{2k} \, x^k \\ ~~ g_1(x) := c_1 + c_3 x + c_5 x^2 + \cdots = \sum_{0 \leq k < n/2} c_{2k+1} \, x^k $$
とすると、$g_0, g_1 \in \mathbf{C}^{~n/2}$であり、
$$ ~~ f(x) = g_0(x^2) + x g_1(x^2) $$
と表せます。

次に、$\mathcal{F}_n(f) = (f(\omega_n^{~k}))$のそれぞれの成分を求めていくと、$0 \leq k < n/2$に対して、
$$ ~~ f(\omega_n^{~k}) = g_0(\omega_{n/2}^{~k}) + \omega_n^{~k} \, g_1(\omega_{n/2}^{~k}) \\ ~~ f(\omega_n^{~k+n/2}) = f(-\omega_n^{~k}) = g_0(\omega_{n/2}^{~k}) - \omega_n^{~k} \, g_1(\omega_{n/2}^{~k}) $$
と計算できます。
つまり、$\mathcal{F}_n(f)$を前後半に分割して考えると、$\mathcal{F}_{n/2}(g_0) = (g_0(\omega_{n/2}^{~k}))$および$\mathcal{F}_{n/2}(g_1) = (g_1(\omega_{n/2}^{~k}))$が既知であれば、$\mathcal{F}_n(f)$を時間計算量$O(n)$で求められます。

上の式を書き換えて、
$$ ~~ \mathcal{F}_n(f)_k = \mathcal{F}_{n/2}(g_0)_k + \omega_n^{~k} \, \mathcal{F}_{n/2}(g_1)_k \\ ~~ \mathcal{F}_n(f)_{k+n/2} = \mathcal{F}_{n/2}(g_0)_k - \omega_n^{~k} \, \mathcal{F}_{n/2}(g_1)_k $$
のようにも表現できます。

また、$n=1$のときは$\mathcal{F}_1(f) = c_0 \in \mathbf{C}$とすればよいです。
時間計算量については、上記の操作を$n$が$2$の冪のときだけ実行するため、$O(n \log n)$となります。

以上の考察を定理の形でまとめておきます。

このように、サイズの小さい同じ問題に分割して再帰的に解く方法は分割統治法 (divide-and-conquer algorithm) と呼ばれます。例えば、マージソートも同様の手法によります。

FFT の逆変換である逆高速フーリエ変換 (inverse fast Fourier transform, IFFT) についても、同様の考察から次の漸化式が得られます。

FFT と異なる点は、$\omega_n^{~k}$の代わりに$\omega_n^{-k}$を使うことと、各階層ごとに計算結果を$2$で割ることです。

また、次のように DFT の結果を利用して IDFT の結果を求める方法もあります。
これについては FFT である必要はなく、前回までの議論でも成り立ちます。

つまり、DFT の結果に対して、$0< k< n$の範囲を逆順に並べ替え、$n$で割れば IDFT の結果になるということです。

以上で、DFT および IDFT が高速化され、時間計算量$O(n \log n)$で実行できることがわかりました。
これにより、畳み込み (多項式の積) も$O(n \log n)$に高速化されます。

プログラミング

定理 1 および定理 3 で得られた式をそのまま実装します。
定理 1 の主張が再帰で表現されているため、この実装においても再帰関数を使います。

入力の配列の長さが$2$の冪でない場合、それを超える最小の$2$の冪を$n$とします。戻り値の配列の長さもその$n$となります。
Convolution 関数については前回の実装とほぼ同じで、$n$を$2$の冪に変更する部分のみ異なります。

      using System;
using System.Numerics;

public static class FFT
{
    public static int ToPowerOf2(int n)
    {
        var p = 1;
        while (p < n) p <<= 1;
        return p;
    }

    // k 番目の 1 の n 乗根 (n-th root)
    static Complex NthRoot(int n, int k)
    {
        var t = 2 * Math.PI * k / n;
        return Complex.FromPolarCoordinates(1, t);
    }

    // c の長さは 2 の冪とします。
    static void TransformRecursive(Complex[] c)
    {
        var n = c.Length;
        if (n == 1) return;

        var n2 = n >> 1;
        var c0 = new Complex[n2];
        var c1 = new Complex[n2];
        for (int k = 0; k < n2; ++k)
        {
            c0[k] = c[2 * k];
            c1[k] = c[2 * k + 1];
        }

        TransformRecursive(c0);
        TransformRecursive(c1);

        for (int k = 0; k < n2; ++k)
        {
            var v0 = c0[k];
            var v1 = c1[k] * NthRoot(n, k);
            c[k] = v0 + v1;
            c[k + n2] = v0 - v1;
        }
    }

    // 戻り値の長さは 2 の冪となります。
    public static Complex[] Transform(Complex[] c, bool inverse)
    {
        if (c == null) throw new ArgumentNullException(nameof(c));

        var n = ToPowerOf2(c.Length);
        var t = new Complex[n];
        c.CopyTo(t, 0);

        TransformRecursive(t);

        if (inverse && n > 1)
        {
            Array.Reverse(t, 1, n - 1);
            for (int k = 0; k < n; ++k) t[k] /= n;
        }
        return t;
    }

    // 戻り値の長さは 2 の冪となります。
    public static Complex[] Convolution(Complex[] a, Complex[] b)
    {
        if (a == null) throw new ArgumentNullException(nameof(a));
        if (b == null) throw new ArgumentNullException(nameof(b));

        var n = ToPowerOf2(a.Length + b.Length - 1);
        Array.Resize(ref a, n);
        Array.Resize(ref b, n);

        var fa = Transform(a, false);
        var fb = Transform(b, false);

        for (int k = 0; k < n; ++k)
        {
            fa[k] *= fb[k];
        }
        return Transform(fa, true);
    }
}
    

逆変換のみに必要な処理 (定理 3) は、Transform 関数の最後に追加しています。

以下のオンラインジャッジで FFT のコードをテストできます。時間計算量$O(n^2)$のナイーブ解では通らず、$O(n \log n)$の解で通ります。

• 高速フーリエ変換 (AtCoder)
• Multiplication of Big Integers II (AOJ)

作成したサンプル コード

• DftNttTest (GitHub)

以上で、再帰関数により FFT を実装できました。数式で表現した通りの実装をすればよいため仕組みがわかりやすいという利点があるものの、この実装では各層で配列を作成して値をコピーしているため、空間計算量も$O(n \log n)$となります。
次回はビット反転置換によりこれを改善します。

前回: (2) 数列の畳み込み
次回: (4) ビット反転置換