大整数数组快速卷积方案：Scipy FFT卷积结果错误问题

我需要对长整数数组执行卷积运算，数组内数值较大但仍在float64范围内。使用Scipy的convolve函数FFT方法时得到错误结果，但该方法速度很快；改用direct方法能得到正确结果，但速度慢很多。输入数组始终由整数组成。

最小复现示例 #

from scipy.signal import convolve
import numpy as np

A = np.array([0] + [1e100]*10000)
convolve(A, A)

错误输出：

array([ 9.15304445e+187, -7.04080342e+187,  1.00000000e+200, ...,
        3.00000000e+200,  2.00000000e+200,  1.00000000e+200])

FFT方法速度：

%timeit convolve(A, A)
458 µs ± 5.09 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Direct方法结果与速度 #

convolve(A, A, method="direct")

正确输出：

array([0.e+000, 0.e+000, 1.e+200, ..., 3.e+200, 2.e+200, 1.e+200])

Direct方法速度：

%timeit convolve(A, A, method="direct")
23.4 ms ± 511 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

非极端规模下的问题 #

即使输入规模不极端，FFT方法仍会出错：

A = np.array([0] + [1e10]*10000)
convolve(A, A)

错误输出：

array([1.96826156e+08, 1.35742176e+08, 1.00000000e+20, ...,
       3.00000000e+20, 2.00000000e+20, 1.00000000e+20])

悬赏针对示例：A = np.array([0] + [1e10]*10000)

问题根源是FFT运算引入的浮点精度误差：当输入数值量级差异大（如包含0和超大整数）时，FFT的舍入误差会被放大，导致结果偏离正确值。以下是几种快速获取正确结果的方案：

1. 针对固定模式数组直接生成结果 #

如果你的数组存在固定模式（如开头为0，后续为相同大整数），无需调用卷积函数，直接根据卷积的滑动求和逻辑生成结果，这是速度最快的方案：

import numpy as np

A = np.array([0] + [1e10]*10000)
val = (A[1]) ** 2
n = len(A)
# 构造结果数组：前两位为0，中间从val递增到(n-1)*val，再递减回val
result = np.concatenate([
    np.zeros(2),
    np.arange(1, n) * val,
    np.arange(n-2, 0, -1) * val
])

该方法时间复杂度为O(n)，远快于FFT和direct卷积。

2. 用Numba加速Direct卷积 #

通用场景下，使用Numba对direct卷积逻辑进行JIT编译，可大幅提升速度，接近FFT方法的同时保证结果准确：

import numba
import numpy as np

@numba.jit(nopython=True)
def fast_convolve(a, b):
    n, m = len(a), len(b)
    result = np.zeros(n + m - 1, dtype=a.dtype)
    for i in range(n):
        if a[i] == 0:
            continue  # 跳过0元素，进一步提速
        for j in range(m):
            result[i+j] += a[i] * b[j]
    return result

A = np.array([0] + [1e10]*10000, dtype=np.float64)
result = fast_convolve(A, A)

测试显示，该方法速度比Scipy的direct卷积快一个数量级，且结果完全正确。

3. 使用整数类型数组运算 #

由于输入均为整数，可直接使用numpy的整数类型（如int64、int128，需根据数值范围选择），避免浮点精度损失。若数值超出标准整数类型范围，可使用object类型存储Python原生大整数：

A = np.array([0] + [10**10]*10000, dtype=np.int64)
result = np.convolve(A, A)  # numpy默认direct方法，结合Numba可进一步提速

4. 高浮点精度FFT（效果有限） #

若坚持使用FFT方法，可尝试float128类型提升精度，但会牺牲部分速度，且极端数值场景下仍可能存在误差：

A = np.array([0] + [1e10]*10000, dtype=np.float128)
result = convolve(A, A, method='fft')

• 数组有固定模式时，直接生成结果是最优解，兼顾速度与准确性。
• 通用场景优先选择Numba加速的direct卷积，平衡速度与精度。
• 大数值、量级差异大的整数数组，避免使用FFT方法做卷积，浮点精度误差无法避免。

内容的提问来源于stack exchange，提问作者Simd