基于FFT多项式乘法判断a、b元素和是否存在于集合c的问题

解法逻辑 #

核心原理 #

当我们把集合a、b转换为多项式时，规则为：若数值k属于集合，则多项式x^{k项的系数为1，否则为0。两个多项式相乘后，结果中x}m项的系数，就是a和b中元素相加等于m的总对数，只要系数大于0就说明存在满足n1+n2=m的n1∈a、n2∈b。

现有代码错误修正 #

你当前构建系数数组的逻辑完全错误，不是按集合元素顺序依次填充数组前几位，而是要按元素值对应下标设置系数为1：

# 正确的系数构建逻辑
a_coeffs = [0]*(n+1)
b_coeffs = [0]*(n+1)
for num in a:
    a_coeffs[num] = 1
for num in b:
    b_coeffs[num] = 1

结果判断步骤 #

得到c_coeffs后：

• 由于FFT计算存在浮点精度误差，不能直接判断系数是否大于0，需要设置极小阈值（如1e-6）过滤误差
• 遍历集合c中的每个元素s，只要存在s对应的c_coeffs[s]大于阈值，就返回True，否则返回False

修正后的完整代码 #

from numpy.fft import fft, ifft
from numpy import real

def polynomial_multiply(a_coeff_list, b_coeff_list):
    len_a = len(a_coeff_list)
    len_b = len(b_coeff_list)
    # 补0到足够长度避免循环卷积
    required_len = len_a + len_b - 1
    a_coeff_list = a_coeff_list + [0]*(required_len - len_a)
    b_coeff_list = b_coeff_list + [0]*(required_len - len_b)
    a_fft = fft(a_coeff_list)
    b_fft = fft(b_coeff_list)
    c_fft = a_fft * b_fft
    inverse_c = ifft(c_fft)
    return real(inverse_c)

def check_sum_exists(a, b, c, n):
    a_coeffs = [0]*(n+1)
    b_coeffs = [0]*(n+1) 
    # 正确构建系数数组
    for num in a:
        a_coeffs[num] = 1
    for num in b:
        b_coeffs[num] = 1
    # 多项式乘法
    c_coeffs = polynomial_multiply(a_coeffs, b_coeffs)
    # 阈值过滤浮点误差
    eps = 1e-6
    for s in c:
        if s < len(c_coeffs) and c_coeffs[s] > eps:
            return True
    return False

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