每日算法 — 使用java实现NTT数论变换：模意义下多项式乘法与卷积加速

引言

快速傅里叶变换（FFT）通过将多项式乘法从O(n²)加速到O(n log n)，在信号处理、大整数运算等领域发挥着重要作用。然而，FFT基于浮点数运算，存在精度误差问题。数论变换（Number Theoretic Transform，NTT）是FFT在模运算域上的精确实现，利用原根替代单位复根，完全避免了浮点精度问题，在算法竞赛和密码学中有广泛应用。

算法原理

1. 从FFT到NTT

FFT的核心是利用单位复根的性质：
– ωₙ = e^(2πi/n) 是n次单位复根
– 满足 ωₙⁿ = 1 且 ωₙ^k (k=0..n-1) 互不相同

NTT将这一思想迁移到模p的整数域中：
– 寻找模p的原根g，使得 g^((p-1)/n) 是n次单位根
– 要求 n | (p-1)，即n整除p-1

2. 常用模数与原根

3. 核心性质

• 原根存在性：当且仅当 p = 2, 4, q^k, 2q^k（q为奇素数）时，模p存在原根
• 蝴蝶操作：与FFT相同的分治结构，只是运算在模p下进行
• 逆变换：利用原根的逆元实现，同样需要满足特定条件

Java实现

/**
 * NTT数论变换实现
 * 模数：998244353 = 119 * 2^23 + 1，原根为3
 * 支持长度不超过2^23的多项式乘法
 */
public class NTT {

    // 常用模数：998244353 = 119 * 2^23 + 1
    static final int MOD = 998244353;
    static final int G = 3;  // 原根

    /**
     * 快速幂取模
     */
    static long power(long base, long exp, int mod) {
        long result = 1;
        base %= mod;
        while (exp > 0) {
            if ((exp & 1) == 1) result = result * base % mod;
            base = base * base % mod;
            exp >>= 1;
        }
        return result;
    }

    /**
     * 模逆元（费马小定理，要求mod为素数）
     */
    static long inv(long a, int mod) {
        return power(a, mod - 2, mod);
    }

    /**
     * NTT核心变换
     * @param a 多项式系数数组（会被修改）
     * @param invert true表示逆变换
     */
    static void ntt(long[] a, boolean invert) {
        int n = a.length;

        // 位逆序置换（蝴蝶操作的前提）
        for (int i = 1, j = 0; i < n; i++) {
            int bit = n >> 1;
            for (; j >= bit; bit >>= 1) {
                j -= bit;
            }
            j += bit;
            if (i < j) {
                long temp = a[i];
                a[i] = a[j];
                a[j] = temp;
            }
        }

        // 逐层进行NTT
        for (int len = 2; len <= n; len <<= 1) {
            // wlen = g^((mod-1)/len) 或它的逆元
            long wlen = power(G, (MOD - 1) / len, MOD);
            if (invert) {
                wlen = inv(wlen, MOD);
            }

            for (int i = 0; i < n; i += len) {
                long w = 1;
                for (int j = 0; j < len / 2; j++) {
                    long u = a[i + j];
                    long v = a[i + j + len / 2] * w % MOD;

                    a[i + j] = (u + v) % MOD;
                    a[i + j + len / 2] = (u - v + MOD) % MOD;

                    w = w * wlen % MOD;
                }
            }
        }

        // 逆变换需要除以n（即乘以n的逆元）
        if (invert) {
            long nInv = inv(n, MOD);
            for (int i = 0; i < n; i++) {
                a[i] = a[i] * nInv % MOD;
            }
        }
    }

    /**
     * 多项式乘法（卷积）
     * @param a 多项式A的系数
     * @param b 多项式B的系数
     * @return 多项式A*B的系数（模998244353）
     */
    static long[] multiply(long[] a, long[] b) {
        if (a.length == 0 || b.length == 0) return new long[0];

        // 找到不小于结果长度的最小2的幂
        int resultSize = a.length + b.length - 1;
        int n = 1;
        while (n < resultSize) n <<= 1;

        // 填充到长度n
        long[] fa = new long[n];
        long[] fb = new long[n];
        System.arraycopy(a, 0, fa, 0, a.length);
        System.arraycopy(b, 0, fb, 0, b.length);

        // 正变换
        ntt(fa, false);
        ntt(fb, false);

        // 点值相乘
        for (int i = 0; i < n; i++) {
            fa[i] = fa[i] * fb[i] % MOD;
        }

        // 逆变换
        ntt(fa, true);

        // 截取有效部分
        long[] result = new long[resultSize];
        System.arraycopy(fa, 0, result, 0, resultSize);
        return result;
    }

    /**
     * 大整数乘法（使用NTT加速）
     * 将大整数按位拆分为多项式系数，相乘后再处理进位
     */
    static String bigIntegerMultiply(String num1, String num2) {
        if (num1.equals("0") || num2.equals("0")) return "0";

        // 转换为数字数组（低位在前）
        int n1 = num1.length(), n2 = num2.length();
        long[] a = new long[n1];
        long[] b = new long[n2];

        for (int i = 0; i < n1; i++) {
            a[n1 - 1 - i] = num1.charAt(i) - '0';
        }
        for (int i = 0; i < n2; i++) {
            b[n2 - 1 - i] = num2.charAt(i) - '0';
        }

        // NTT卷积
        long[] c = multiply(a, b);

        // 处理进位
        long carry = 0;
        for (int i = 0; i < c.length; i++) {
            long sum = c[i] + carry;
            c[i] = sum % 10;
            carry = sum / 10;
        }

        // 转换为字符串
        StringBuilder sb = new StringBuilder();
        int i = c.length - 1;
        while (i >= 0 && c[i] == 0) i--;
        while (i >= 0) {
            sb.append(c[i]);
            i--;
        }

        return sb.length() == 0 ? "0" : sb.toString();
    }

    /**
     * 序列卷积（用于组合计数等场景）
     */
    static long[] convolution(long[] a, long[] b) {
        return multiply(a, b);
    }

    public static void main(String[] args) {
        // 测试1：多项式乘法
        // (1 + 2x + 3x^2) * (4 + 5x) = 4 + 13x + 22x^2 + 15x^3
        long[] p1 = {1, 2, 3};
        long[] p2 = {4, 5};
        long[] result = multiply(p1, p2);
        System.out.print("多项式乘法结果：");
        for (long x : result) System.out.print(x + " ");
        System.out.println(); // 应输出：4 13 22 15

        // 测试2：大整数乘法
        String num1 = "123456789";
        String num2 = "987654321";
        String product = bigIntegerMultiply(num1, num2);
        System.out.println(num1 + " * " + num2 + " = " + product);
        System.out.println("验证：" + (product.equals("121932631112635269") ? "正确" : "错误"));

        // 测试3：大数乘法
        String big1 = "12345678901234567890";
        String big2 = "98765432109876543210";
        String bigProduct = bigIntegerMultiply(big1, big2);
        System.out.println("大数乘法结果长度：" + bigProduct.length());
    }
}

复杂度分析

相比朴素的O(n²)算法，NTT在处理大规模数据时优势显著。

应用场景

1. 大整数运算：密码学中的RSA、ECC等需要高精度乘法
2. 组合计数：生成函数相乘计算组合数
3. 信号处理：离散卷积的精确计算
4. 字符串匹配：基于多项式的模式匹配算法

总结

NTT是FFT在有限域上的精确实现，消除了浮点误差的同时保持了O(n log n)的高效性。理解NTT的关键在于掌握原根在单位根中的替代作用，以及模运算下的蝴蝶操作。对于需要精确结果的卷积计算，NTT是不二之选。