文献あり

離散フーリエ変換と数論変換 (5) 数論変換 (NTT)

1367

この連載について

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

本文

前回までに考察した複素数体$\mathbf{C}$上の離散フーリエ変換 (DFT) では、入力が整数であっても小数の演算が必要でした。
そこで今回は、入力が整数の場合には整数の演算のみで済むよう、何らかの整数の集合上でも前回までの理論が成り立たないかを検討します。
この結果として、数論変換 (number-theoretic transform, NTT) が得られます。

実用性を考慮すると、結論としては素数$p$に対して剰余環$\mathbf{Z}/p \mathbf{Z}$上で考えるのが簡単であり、それを前提とした文献も多いですが、本稿では十分条件をなるべく広く取って考察します。

準備

今後の議論に必要な定義および定理を準備します。定理の証明は省略します。

位数・$1$の原始$k$乗根

$a$と$n$が互いに素であるとき、
$a^k \equiv 1 \pmod{n}$を満たす最小の$k>0$を、$a$の$\mathrm{mod} ~n$における位数 (order) と呼ぶ。

このときの$a$、つまり位数$k$の元を、$1$の原始$k$乗根 (primitive $k$-th root of unity) と呼ぶ。

逆元の存在

$a$と$n$が互いに素
$\Longleftrightarrow ~ a$の$\mathrm{mod} ~n$における逆元、すなわち$ax \equiv 1 \pmod{n}$を満たす$x$が存在する

数論変換

加法および乗法が定義されている何らかの整数の集合上で、これまでの DFT の議論で現れた2つの定理 (IDFT および畳み込み) が成り立つための十分条件を考察します。

まず、$\omega$が$1$の原始$n$乗根であるとし、 (1) 離散フーリエ変換 (DFT) と同様に DFT を定義します。
すなわち、$\omega$により生成される$\{ \omega^{~k} ~|~ 0 \leq k < n \}$を標本点の集合とします。
このとき、$\omega^{~k} \neq 1 ~ (0 < k < n), ~~ \omega^{~n} = 1$が成り立っています。

逆離散フーリエ変換 (IDFT) に関する定理が成り立つには、以前の証明の内容から、

(i) $ \displaystyle \sum_{0 \leq k < n} \omega^{~jk} = \left\{ \begin{array}{l} n ~~ (j = 0) \\ 0 ~~ (0 < j < n) \end{array} \right. $
(ii) 逆元$\omega^{-1}$が存在する
(iii) 逆元$n^{-1}$が存在する

が満たされていれば十分であることがわかります。

(i) については、$j=0$のときはつねに成り立つため、$0 < j < n$の場合のみ考えればよいです。
(ii) については、$\omega \cdot \omega^{n-1} = 1$より、$\omega^{-1} = \omega^{n-1}$とすればよいため満たされます。
(iii) については定理 1 により、剰余環$\mathbf{Z}/m \mathbf{Z}$で考えた場合に、$n$と$m$が互いに素であることと同値です。

また、畳み込み定理については、加法および乗法が定義されていれば成り立ちます。
以上の考察から、次のようにまとめられます。

数論変換

$n$と$m$は互いに素であるとする。
$\omega$を剰余環$\mathbf{Z}/m \mathbf{Z}$における$1$の原始$n$乗根の一つとする。
さらに、次が成り立つとする。
$$ ~~ \sum_{0 \leq k < n} \omega^{~jk} \equiv 0 \pmod{m} ~~~ (0 < j < n) $$

このとき、IDFT に関する定理および畳み込み定理は、値の集合と標本点の組$\langle \mathbf{C}, \omega_n \rangle$を$\langle \mathbf{Z}/m \mathbf{Z}, \omega \rangle$に置き換えても成り立つ。
このような剰余環$\mathbf{Z}/m \mathbf{Z}$上の DFT を、数論変換 (number-theoretic transform, NTT) と呼ぶ。

ここまでの議論においては、$n$が$2$の冪である必要も、$m$が素数である必要もありません。
これらの強い制限は、次回の高速化で活用されることになります。

$n=5,6$の例を考えると、次のようにいくつかのパターンがあることがわかります。
$1$の原始$n$乗根$\omega$が存在しても NTT が適用不可であることがあります。

$n=5:$

$\mathbf{Z}/m \mathbf{Z}$	$\omega$	$\{ \omega^k \}$	$\sum \omega^{~jk} \equiv 0$	$n^{-1}$	適用	備考
$\mathbf{Z}/31 \mathbf{Z}$	$2$	$\{ 1, 2, 4, 8, 16 \}$	成立	$25$	可	$m$は素数
$\mathbf{Z}/25 \mathbf{Z}$	$6$	$\{ 1, 6, 11, 16, 21 \}$	$\sum \omega^{~k} \neq 0$	なし	不可
$\mathbf{Z}/55 \mathbf{Z}$	$16$	$\{ 1, 16, 36, 26, 31 \}$	成立	なし	不可
$\mathbf{Z}/22 \mathbf{Z}$	$3$	$\{ 1, 3, 9, 5, 15 \}$	$\sum \omega^{~k} \neq 0$	$9$	不可
$\mathbf{Z}/121 \mathbf{Z}$	$3$	$\{ 1, 3, 9, 27, 81 \}$	成立	$97$	可	$m$は合成数
$\mathbf{Z}/341 \mathbf{Z}$	$4$	$\{ 1, 4, 16, 64, 256 \}$	成立	$273$	可	$m$は合成数

$n=6:$

$\mathbf{Z}/m \mathbf{Z}$	$\omega$	$\{ \omega^k \}$	$\sum \omega^{~jk} \equiv 0$	$n^{-1}$	適用	備考
$\mathbf{Z}/31 \mathbf{Z}$	$6$	$\{ 1, 6, 5, 30, 25, 26 \}$	成立	$26$	可	$m$は素数
$\mathbf{Z}/21 \mathbf{Z}$	$2$	$\{ 1, 2, 4, 8, 16, 11 \}$	$\sum \omega^{~2k} \neq 0$	なし	不可
$\mathbf{Z}/21 \mathbf{Z}$	$5$	$\{ 1, 5, 4, 20, 16, 17 \}$	成立	なし	不可
$\mathbf{Z}/35 \mathbf{Z}$	$4$	$\{ 1, 4, 16, 29, 11, 9 \}$	$\sum \omega^{~2k} \neq 0$	$6$	不可
$\mathbf{Z}/49 \mathbf{Z}$	$19$	$\{ 1, 19, 18, 48, 30, 31 \}$	成立	$41$	可	$m$は合成数
$\mathbf{Z}/91 \mathbf{Z}$	$10$	$\{ 1, 10, 9, 90, 81, 82 \}$	成立	$76$	可	$m$は合成数

$m$が素数の場合は扱いやすいです。
例えば$m=31$のとき、位数$30$の元 (原始根) が存在します。したがって、$30$の任意の約数を位数とする元が存在します。

$m$が合成数の場合、$\sum \omega^{~jk}$についての十分条件を満たす$\omega$を見つけることが難しいです。
例えば$(n, m, \omega) = (5, 22, 3), (6, 35, 4)$は NTT を適用できそうに見えて、実はこの条件を満たしていません。

もし$\omega^{~j} -1$が$\mathrm{mod} ~m$で逆元を持てば、
$$ ~~ \sum_{0 \leq k < n} \omega^{~jk} \equiv 0 ~~ (0 < j < n) $$
であることは示せます。
そこで、$\sum \omega^{~jk}$についての十分条件を「$0 < j < n$に対して、$\omega^{~j} -1$と$m$は互いに素」に変更してもよいかもしれません。

次回は、$m$を素数に限定することで、高速化と実用性を備えた構築方法を得ることになります。

プログラミング

DFT のコードを基に NTT を実装していきましょう。
今回はパラメーター$n, m, \omega, n^{-1}$をコンストラクターで指定できるようにしておきます。
また、$1$の$n$乗根の集合$\{ \omega^k \}$をあらかじめ計算してキャッシュします。

      using System;

public class NTT
{
    // k 番目の 1 の n 乗根 (ω^k)
    static long[] NthRoots(int n, long m, long w)
    {
        var r = new long[n];
        r[0] = 1;
        for (int k = 1; k < n; ++k)
            r[k] = r[k - 1] * w % m;
        return r;
    }

    int n;
    long m, nInv;
    long[] roots;

    public NTT(int n, long m, long w, long nInv)
    {
        this.n = n;
        this.m = m;
        this.nInv = nInv;
        roots = NthRoots(n, m, w);
    }

    // f(ω^k) の値
    long f(long[] c, int k)
    {
        var r = 0L;
        for (int j = 0; j < c.Length; ++j)
            r = (r + c[j] * roots[k * j % n]) % m;
        return r;
    }

    public long[] Transform(long[] c, bool inverse)
    {
        if (c == null) throw new ArgumentNullException(nameof(c));

        var r = new long[n];
        for (int k = 0; k < n; ++k)
            r[k] = inverse ? f(c, n - k) * nInv % m : f(c, k);
        return r;
    }

    public long[] Convolution(long[] a, long[] b)
    {
        if (a == null) throw new ArgumentNullException(nameof(a));
        if (b == null) throw new ArgumentNullException(nameof(b));

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

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

Transform, Convolution 関数のいずれも、時間計算量は$O(n^2)$です。次回にこれを高速化します。
では Convolution 関数をテストしてみましょう。

      using System;

static class Program
{
    static void Main()
    {
        Test(5, 31, 2, 25);
        Test(5, 121, 3, 97);
        Test(5, 341, 4, 273);

        Test(6, 31, 6, 26);
        Test(6, 49, 19, 41);
        Test(6, 91, 10, 76);
    }

    static void Test(int n, long m, long w, long nInv)
    {
        var a = new long[] { 2, 1, 1 };
        var b = new long[] { m - 1, m - 1, 1 };

        var ntt = new NTT(n, m, w, nInv);
        // { m - 2, m - 3, 0, 0, 1 }
        var c = ntt.Convolution(a, b);
    }
}