文献あり

離散フーリエ変換と数論変換 (4) ビット反転置換

649

この連載について

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

本文

前回は高速フーリエ変換 (FFT) の仕組みとそのプログラミングについて説明しました。
定理で記述した漸化式をそのまま実装すればよいため仕組みがわかりやすいという利点があるものの、この実装では各層で配列を作成して値をコピーしているため、空間計算量も$O(n \log n)$となっていました。
そこで、今回はビット反転置換 (bit-reversal permutation) を利用して、空間計算量を$O(n)$に改善します。なお、この用語は文献 [1] によります。

ビット反転置換

前回の FFT に関する定理において、$f = (c_k) \in \mathbf{C}^{~n}$に対し、再帰的に
$$ ~~ \left\{ \begin{array}{l} g_0 := (c_{2k}) \in \mathbf{C}^{~n/2} \\ g_1 := (c_{2k+1}) \in \mathbf{C}^{~n/2} \end{array} \right. $$
と2つの多項式に分割していきました。
つまり、その階層において偶数番目と奇数番目の値の列に分割する操作であり、
例えば$n=8$のとき、元の$\mathbf{C}^8$におけるインデックスは次の図のように移動します。

!FORMULA[6][36584159][0]のときの、各階層で処理されるインデックス (by GeoGebra) $n=8$のときの、各階層で処理されるインデックス (by GeoGebra)

空間計算量が$O(n \log n)$となっていたのは、$g_0, g_1$を表す配列を各階層で生成していたことが原因でした。
そこで、初めから最下層のインデックスである$0, 4, 2, 6, 1, 5, 3, 7$の順に配列の要素を並べ替えておけば、一時格納用の配列を生成する必要がなくなり、$O(n)$に改善できます。
以下では、この置換を求める方法について考察します。

上の図を観察すると、各インデックスを$2$進法で表したとき、下位ビットから見て$0$が多く連続するほど左に移動することがわかります。これに対して移動前は、上位ビットから見て$0$が多く連続するほど左にあります。
このようなパターンを見ることで、次の定理が得られます。

ビット反転置換

$n$を$2$の冪とする。
$f = (c_k) \in \mathbf{C}^{~n}$に対し、
(i) $n=1$のとき、$\mathcal{P}_1(f) := c_0 \in \mathbf{C}$
(ii) $n>1$のとき、
$$ ~~ \left\{ \begin{array}{l} g_0 := (c_{2k}) \in \mathbf{C}^{~n/2} \\ g_1 := (c_{2k+1}) \in \mathbf{C}^{~n/2} \end{array} \right. $$
とし、
$$ ~~ \left\{ \begin{array}{l} \mathcal{P}_n(f)_k := \mathcal{P}_{n/2}(g_0)_k \\ \mathcal{P}_n(f)_{k+n/2} := \mathcal{P}_{n/2}(g_1)_k \end{array} \right. ~~ (0 \leq k < n/2) $$
とする。

このとき、ある整数を$\log n$桁の$2$進数で表したとき、そのビットを逆順に並べ替えたものに写す置換を$\rho_n$で表すと、$f$の$k$番目の成分は$\mathcal{P}_n$により$\rho_n(k)$番目に移動する。
すなわち、
$$ ~~ \mathcal{P}_n(f)_{\rho_n(k)} = c_k $$
が成り立つ。

この$\rho_n$を、ビット反転置換 (bit-reversal permutation) と呼ぶ。

例えば、

$\rho_8(6) = \rho_8(110_2) = 011_2 = 3$
$\rho_{16}(5) = \rho_{16}(0101_2) = 1010_2 = 10$
$\rho_{32}(1) = \rho_{32}(00001_2) = 10000_2 = 16$

です。

この置換$\rho_n$を実装すると、次のようになります。

          static int[] BitReversal_Original(int n)
    {
        var b = new int[n];
        for (int u = 1, d = n >> 1; u < n; u <<= 1, d >>= 1)
            for (int k = 0; k < n; ++k)
                if ((k & u) != 0)
                    b[k] |= d;
        return b;
    }

それぞれのビットごとに$0 \leq k < n$に対してビットフラグを判定しているため時間計算量は$O(n \log n)$です。
これでも十分なのですが、次のようにさらに$O(n)$に効率化できます。

          static int[] BitReversal(int n)
    {
        var b = new int[n];
        for (int p = 1, d = n >> 1; p < n; p <<= 1, d >>= 1)
            for (int k = 0; k < p; ++k)
                b[k | p] = b[k] | d;
        return b;
    }

これは、$n=8$の例において、

$\{ 0 \}$の各要素に$2^2$を加えて$\{ 4 \}$を得る
$\{ 0, 4 \}$の各要素に$2^1$を加えて$\{ 2, 6 \}$を得る
$\{ 0, 4, 2, 6 \}$の各要素に$2^0$を加えて$\{ 1, 5, 3, 7 \}$を得る

という手順になっています。

FFT への適用

再帰関数による実装

入力のベクトルにビット反転置換を施し、前回に実装した FFT の再帰関数を書き換えてみましょう。
なお、置換前と置換後には対称性があるため、「コピー先のインデックス」と考えても「コピー元のインデックス」と考えても同じです。

          // c の長さは 2 の冪とします。
    // h: 更新対象の長さの半分
    static void TransformRecursive(Complex[] c, int l, int h)
    {
        if (h == 0) return;

        TransformRecursive(c, l, h >> 1);
        TransformRecursive(c, l + h, h >> 1);

        for (int k = 0; k < h; ++k)
        {
            var v0 = c[l + k];
            var v1 = c[l + k + h] * NthRoot(h << 1, k);
            c[l + k] = v0 + v1;
            c[l + k + h] = 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 br = BitReversal(n);

        var t = new Complex[n];
        for (int k = 0; k < c.Length; ++k)
            t[br[k]] = c[k];

        TransformRecursive(t, 0, n >> 1);

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

これで空間計算量を$O(n)$に改善でき、さらに、その分だけ実行速度も上がります。

非再帰関数による実装

上記と同等の処理を、再帰のない関数として次のように実装することもできます。
再帰関数をトップダウンで順次呼び出していく方法とは逆に、ループを利用してボトムアップで各階層を処理します。

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

        var n = ToPowerOf2(c.Length);
        var br = BitReversal(n);

        var t = new Complex[n];
        for (int k = 0; k < c.Length; ++k)
            t[br[k]] = c[k];

        for (int h = 1; h < n; h <<= 1)
        {
            for (int l = 0; l < n; l += h << 1)
            {
                for (int k = 0; k < h; ++k)
                {
                    var v0 = t[l + k];
                    var v1 = t[l + k + h] * NthRoot(h << 1, k);
                    t[l + k] = v0 + v1;
                    t[l + k + h] = v0 - v1;
                }
            }
        }

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