文献あり

離散フーリエ変換と数論変換 (6) NTT の高速化

2253

この連載について

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

本文

前回は離散フーリエ変換 (DFT) を剰余環 $Z / m Z$ 上で考える数論変換 (NTT) を説明し、時間計算量 $O (n^{2})$ の実装を紹介しました。
今回はこれに高速フーリエ変換 (FFT) と同様の手法を適用して、 $O (n \log n)$ に高速化します。

用語

まず、NTT についての資料などでたびたび言及される高速剰余変換 (fast modulo transform, FMT) という用語の出典について調べてみたところ、初期の信頼できる文献として次のものが見つかりました。

どうやら本稿の「NTT に FFT の手法を適用して高速化したもの」という文脈とは異なるようです。
また、英語で調べてもこの用語は見当たりませんでした。

次の論文では、「NTT に FFT の手法を適用して高速化したもの」が高速数論変換 (fast number-theoretic transform, FNTT) という用語で定義されています。一般的に使われる用語になっているかどうかはわかりません。

準備

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

原始根・生成元

$p$ を素数とする。
$\mod p$ における位数が $p - 1$ である元を、 $\mod p$ における原始根 (primitive root) または生成元 (generator) と呼ぶ。

原始根または生成元のうちの一つを $g$ で表すことが多いです。
この定義は、 $g^{k}$ により $1$ から $p - 1$ までの数を生成するような $g$ と言い換えられます。
すなわち、
${g^{k} (\mod p) | 0 < k < p} = {x | 0 < x < p}$
を満たす $g$ を指し、こちらを定義とすることもあります。

原始根の存在定理

$p$ を素数とするとき、 $\mod p$ における原始根は存在する。

なお、原始根はただ一つに定まるとは限りません。

位数

d

の元の個数

$p$ を素数とし、 $d$ を $p - 1$ の約数とする。
このとき、 $\mod p$ における位数が $d$ の元の個数は $φ (d)$ である。
ここで、 $φ$ は Euler の totient 関数である。

したがって、 $\mod p$ における原始根の個数は $φ (p - 1)$ です。
なお、この定理から $φ$ に関する性質 $\sum_{d ∣ n} φ (d) = n$ も得られます。

Dirichlet の算術級数定理

$a$ と $d$ は互いに素であるとする。
等差数列 $a_{n} := a + d n (n = 0, 1, 2, \dots)$ は、素数を無限に含む。

$p$ を $3$ 以上の素数とし、 $g$ を $\mod p$ における原始根の一つとするとき、
$g^{\frac{p - 1}{2}} \equiv - 1 (\mod p)$

$x := g^{\frac{p - 1}{2}}$ とし、 $x \neq \pm 1$ であると仮定すると、 $(x + 1) (x - 1) \neq 0$ となるはずである。
しかし、 $(x + 1) (x - 1) = x^{2} - 1 \equiv g^{p - 1} - 1 \equiv 0$ に矛盾するため、 $x$ は $\pm 1$ のいずれかとなる。
$g^{\frac{p - 1}{2}}$ と $g^{p - 1} \equiv 1$ は相異なる数だから、 $g^{\frac{p - 1}{2}} \equiv - 1$ となる。

NTT の高速化

前回に考察した NTT を、FFT と同様の手法により高速化するための十分条件について考察します。
(3) 高速フーリエ変換 (FFT) における考察から、FFT に関する定理が成り立つには、

(i) $n$ が $2$ の冪
(ii) $n > 1$ のとき、 $ω^{2}$ は $1$ の原始 $n / 2$ 乗根
(iii) $n > 1$ のとき、 $ω^{n / 2} = - 1$

が満たされるとよいことがわかります。
なお、(i) が成り立つときは、(ii) も成り立っています。

いったん、ここまでの考察を次の定理にまとめます。

$n$ を $2$ の冪とし、 $n$ と $m$ は互いに素であるとする。
$ω$ を剰余環 $Z / m Z$ における $1$ の原始 $n$ 乗根の一つとする。
次の式が成り立つとする。
$\sum_{0 \leq k < n} ω^{j k} \equiv 0 (\mod m) (0 < j < n)$
また、 $n > 1$ のとき、 $ω^{n / 2} \equiv - 1$ を満たすとする。

このとき、FFT に関する定理は、値の集合と標本点の組 $⟨ C, ω_{n} ⟩$ を $⟨ Z / m Z, ω ⟩$ に置き換えても成り立つ。

$n > 1$ のとき、 $m$ は奇数であると言い換えられます。

$n = 4, 8$ の例を以下に挙げます。
$1$ の原始 $n$ 乗根 $ω$ が存在しても NTT の高速化が適用不可であることがあります。

$n = 4 :$

$Z / m Z$	$ω$	${ω^{k}}$	$\sum ω^{j k} \equiv 0$	適用	備考
$Z / 13 Z$	$5$	${1, 5, - 1, - 5}$	成立	可	$m$ は素数
$Z / 15 Z$	$2$	${1, 2, 4, 8}$	$\sum ω^{2 k} \neq 0$	不可
$Z / 25 Z$	$7$	${1, 7, - 1, - 7}$	成立	可	$m$ は合成数

$n = 8 :$

$Z / m Z$	$ω$	${ω^{k}}$	$\sum ω^{j k} \equiv 0$	適用	備考
$Z / 17 Z$	$2$	${1, 2, 4, 8, - 1, - 2, - 4, - 8}$	成立	可	$m$ は素数
$Z / 51 Z$	$2$	${1, 2, 4, 8, 16, 32, 13, 26}$	$\sum ω^{2 k} \neq 0$	不可
$Z / 289 Z$	$110$	${1, 110, 251, 155, - 1, - 110, - 251, - 155}$	成立	可	$m$ は合成数

しかし、このような $(n, m, ω)$ の組を見つける方法はアドホックになってしまいます。
特定の位数 $n$ を持つ元 $ω$ を見つけるには、原始根から求める方法が簡単です。
そこで、実用における利便性のために強めの制約として、 $m$ を素数 $p$ とすることを考えます。
これにより、次のように定理の十分条件を整理できます。

原始根の存在定理により、 $Z / p Z$ における原始根が存在するため、その一つを $g$ とします。
$g$ の位数は $p - 1$ であることから、 $1$ の原始 $n$ 乗根が存在するには $n ∣ p - 1$ が成り立つように $p$ を選んでおく必要があります。
$d n = p - 1 (d > 0)$ とするとき、 $ω := g^{d}$ と決めることができます。
この $p = d n + 1 (d > 0)$ と表せる素数 $p$ の存在は、Dirichlet の算術級数定理により保証されます。

このとき、定理 4 より、 $n > 1$ に対して $ω^{n / 2} \equiv - 1$ が成り立ちます。
さらに、 $0 < j < n$ に対して、 $ω^{j} - 1$ と $p$ は互いに素であることから、
$\sum_{0 \leq k < n} ω^{j k} \equiv 0 (\mod p)$
が成り立ちます。

以上の考察から、次のようにまとめられます。

$n$ を $2$ の冪、 $p$ を素数とする。
ある $d > 0$ が存在し、 $p = d n + 1$ を満たすとする。

このとき、 $Z / p Z$ における原始根の一つを $g$ とし、 $ω := g^{d}$ とすると、FFT に関する定理は、値の集合と標本点の組 $⟨ C, ω_{n} ⟩$ を $⟨ Z / p Z, ω ⟩$ に置き換えても成り立つ。

前の定理における $m$ を素数と決めることで原始根 $g$ の存在を保証し、 $ω$ を副次的に決めることができるため実用しやすいものとなっています。

$⟨ C, ω_{n} ⟩$ に対しては、いくらでも大きい $n$ を利用することができましたが、 $⟨ Z / p Z, ω ⟩$ は $p$ にも依存しています。したがって、多項式の次数が高くなるほど、大きい $p$ を選ぶ必要があります。
また、 $C$ 上で考える場合は値の制限を気にする必要がありませんでしたが、 $Z / p Z$ においては剰余類の集合 $[0, p)$ に制限されるため、入出力となる値が大きいほど、大きい $p$ を選ぶ必要があります。

また、ある $n$ で定理が成り立つ場合、 $n / 2$ としても成り立つため、任意の $n - 1$ 次以下の多項式に適用できます。
したがって、実用においては $n, p$ をあらかじめ大きな値で構成しておけば変更することなく再利用できます。
とくに、 $p - 1$ が素因数として $2$ を多く含むとよいです。

例として、
$p = 998244353 = 7 \cdot 17 \cdot 2^{23} + 1$
は $n = 2^{18}$ としても $n = 2^{23}$ としても利用できますが、計算する値が $p$ 以上となる場合 (10億など) は $\mod p$ における値となります。

通常は、 $p, g$ を決めて NTT の高速化を利用することになるでしょう。
$p = d n + 1 (d > 0)$ を満たす $p, d$ の組、そして原始根 $g$ はプログラムを作るなどして求めることができます。
原始根 $g$ については、Wolfram|Alpha で PrimitiveRoot[998244353] のように調べる方法もあります。
特徴的な $p, d, g$ の組を求めた結果をいくつか挙げておきます。

$p$	$d n + 1$	$g$	備考
$65537$	$1 \cdot 2^{16} + 1$	$3$
$167772161$	$5 \cdot 2^{25} + 1$	$3$
$998244353$	$119 \cdot 2^{23} + 1$	$3$	一部の分野でよく利用される
$1107296257$	$33 \cdot 2^{25} + 1$	$10$	$10^{9} < p < 2^{31}$
$2013265921$	$15 \cdot 2^{27} + 1$	$31$	$2 \cdot 10^{9} < p < 2^{31}$

プログラミング

(4) ビット反転置換と (5) 数論変換 (NTT) のコードを組み合わせます。
パラメーター $n, p, g$ をコンストラクターで指定できるようにしておきます。

      using System;

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

    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;
    }

    long MPow(long b, long i)
    {
        long r = 1;
        for (; i != 0; b = b * b % p, i >>= 1) if ((i & 1) != 0) r = r * b % p;
        return r;
    }

    // k 番目の 1 の n 乗根 (0 <= k < n/2)
    long[] NthRoots(int n, long w)
    {
        var r = new long[n >> 1];
        r[0] = 1;
        for (int k = 1; k < r.Length; ++k)
            r[k] = r[k - 1] * w % p;
        return r;
    }

    int n;
    public int Length => n;
    long p, nInv;
    int[] br;
    long[] roots;

    // length は 2 の冪に変更されます。
    public FNTT(int length, long p = 998244353, long g = 3)
    {
        n = ToPowerOf2(length);
        this.p = p;
        nInv = MPow(n, p - 2);
        br = BitReversal(n);
        roots = NthRoots(n, MPow(g, (p - 1) / n));
    }

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

        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] * roots[d * k] % p;
            c[l + k] = (v0 + v1) % p;
            c[l + k + h] = (v0 - v1 + p) % p;
        }
    }

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

        var t = new long[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] = t[k] * nInv % p;
        }
        return t;
    }

    // 戻り値の長さは 2 の冪となります。
    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] % p;
        }
        return Transform(fa, true);
    }
}

これで NTT を $O (n \log n)$ に高速化できました。

使用例は次のようになります。
$p, g$ を変更して利用したい場合はコンストラクターで指定できます。

      using System;

static class Program
{
    static void Main()
    {
        var a = new long[] { 1, 2, 3, 4 };
        var b = new long[] { 5, 6, 7, 8, 9 };
        var n = a.Length + b.Length - 1;

        var fntt = new FNTT(n);
        var c = fntt.Convolution(a, b);
        // { 5, 16, 34, 60, 70, 70, 59, 36 }

        var fntt2 = new FNTT(n, 1107296257, 10);
        var c2 = fntt2.Convolution(a, b);
        // { 5, 16, 34, 60, 70, 70, 59, 36 }
    }
}