離散フーリエ変換と数論変換 (7) 中国剰余定理による値の復元

この連載について

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

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

本文

前回は数論変換 (NTT) に高速フーリエ変換 (FFT) の手法を適用しました。
しかし$\mathbf{Z}/p \mathbf{Z}$上で計算すると、値の範囲は$[0, p)$に制限されてしまいます。
そこで今回は、求めたい値が$p$以上になる場合の値の復元方法について説明します。

準備

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

中国剰余定理による値の復元

中国剰余定理 (Chinese remainder theorem) の性質を利用して、NTT の計算についても2つの素数$p_1,p_2$に対し、$\mathbf{Z}/p_1 \mathbf{Z}, \mathbf{Z}/p_2 \mathbf{Z}$上で計算した結果から$\mathbf{Z}/p_1 p_2 \mathbf{Z}$上での値が得られます。すなわち、$[0, p_1 p_2)$の範囲で正確な整数値を求めることができます。
3つ以上の素数でも同様で、例えば$p_1, p_2, p_3$に対して、$p_1 p_2$と$p_3$は互いに素であるため$\mathbf{Z}/p_1 p_2 p_3 \mathbf{Z}$上での値が得られます。

中国剰余定理を表すコードは C# で次のように書けます。

      using System;
using System.Numerics;

public class CRT
{
    // ax + by = 1 の解
    // 前提: a と b は互いに素
    public static (long x, long y) ExtendedEuclid(long a, long b)
    {
        if (b == 0) throw new ArgumentOutOfRangeException(nameof(b));
        if (b == 1) return (1, 1 - a);

        var q = Math.DivRem(a, b, out var r);
        var (u, v) = ExtendedEuclid(b, r);
        return (v, u - q * v);
    }

    static BigInteger MInt(BigInteger x, long mod) => (x %= mod) < 0 ? x + mod : x;

    long mn;
    BigInteger mu, nv;

    // 前提: m と n は互いに素
    public CRT(long m, long n)
    {
        mn = m * n;
        (BigInteger u, BigInteger v) = ExtendedEuclid(m, n);
        mu = MInt(m * u, mn);
        nv = MInt(n * v, mn);
    }

    // a (mod m) かつ b (mod n) である値 (mod mn で一意)
    public long Solve(long a, long b)
    {
        var x = a * nv + b * mu;
        return (long)(x % mn);
    }
}
    

このコードでは、$mn$が 64 ビット整数の範囲に収まる想定で long 型としています。
$mn$が 64 ビット整数の範囲を超える場合には、初めからすべての型を BigInteger にしておくとよいでしょう。

与えられる条件が3種類以上になっても適用できます。
例えば、「孫子算経」の「3で割ると2余り、5で割ると3余り、7で割ると2余る数」は次のように求められます。

      using System;

static class Program
{
    static void Main()
    {
        var crt1 = new CRT(3, 5);
        var crt2 = new CRT(3 * 5, 7);

        var x1 = crt1.Solve(2, 3);
        var x2 = crt2.Solve(x1, 2); // 23
    }
}
    

では、これと前回の FNTT クラスを利用して、巨大な係数を持つ多項式の積を求めてみましょう。

      using System;
using System.Linq;

static class Program
{
    static void Main()
    {
        const long p1 = 998244353, g1 = 3;
        const long p2 = 1107296257, g2 = 10;

        var fntt1 = new FNTT(8, p1, g1);
        var fntt2 = new FNTT(8, p2, g2);

        var a = new long[] { 1000000, 1000000, 1000000 };
        var b = new long[] { 1000000, 1000000, 1000000, 1000000 };

        var c1 = fntt1.Convolution(a, b);
        var c2 = fntt2.Convolution(a, b);

        var crt = new CRT(p1, p2);
        var c = c1.Zip(c2, (x, y) => crt.Solve(x, y)).ToArray();
        // { 1000000000000, 2000000000000, 3000000000000, 3000000000000, 2000000000000, 1000000000000, 0, 0 }
    }
}
    

中国剰余定理は、NTT に限らず、組合せ爆発するような場合の数 (数え上げ) を求める用途などでも利用できます。

作成したサンプル コード

• DftNttTest (GitHub)

前回: (6) NTT の高速化