Python中有没有比numpy polyfromroots更快的批量生成多项式方法？

问题核心原因 #

原生np.polynomial.polynomial.polyfromroots是为单组根设计的接口，循环调用40000次会产生大量Python层的函数调用与循环开销，是耗时的核心来源。

可行加速方案 #

方案1：固定次数多项式向量化展开（4根场景下最优） #

因为你每组都是固定4个根，四次首一多项式的系数可以直接通过初等对称多项式推导得到，完全不需要循环调用接口，直接用numpy向量化运算一次性算出所有组的系数，性能提升可达几十倍。
四次首一多项式（升幂排列，和np.polyfromroots返回顺序一致）系数计算公式如下：
设4个根为$r_1, r_2, r_3, r_4$，升幂系数为$[c_0, c_1, c_2, c_3, c_4]$，则：
$c_4 = 1$
$c_3 = -(r_1 + r_2 + r_3 + r_4)$
$c_2 = r_1r_2 + r_1r_3 + r_1r_4 + r_2r_3 + r_2r_4 + r_3r_4$
$c_1 = -(r_1r_2r_3 + r_1r_2r_4 + r_1r_3r_4 + r_2r_3r_4)$
$c_0 = r_1r_2r_3r_4$

对应代码示例：

import numpy as np

# roots为形状(40000, 4)的二维数组，每行对应一组根
roots = np.random.rand(40000, 4) 

# 向量化计算所有系数
sum_r = roots.sum(axis=1)
sum_r2 = (roots ** 2).sum(axis=1)
sum_rr = (roots[:, None] @ roots[..., None]).ravel() - sum_r2
sum_rrr = roots.prod(axis=1) * (1 / roots).sum(axis=1)
prod_r = roots.prod(axis=1)

# 输出升幂排列的系数，和np.polyfromroots返回格式完全一致
coeffs = np.stack([prod_r, -sum_rrr, sum_rr, -sum_r, np.ones(roots.shape[0])], axis=1)

方案2：Numba JIT编译（适合根次数不固定的场景） #

如果你的每组根数量不固定，可以用Numba将批量处理的逻辑编译为机器码，消除Python循环开销。示例代码如下：

import numpy as np
from numba import jit

@jit(nopython=True)
def batch_polyfromroots(roots_batch):
    n_samples = roots_batch.shape[0]
    n_roots = roots_batch.shape[1]
    coeffs = np.zeros((n_samples, n_roots + 1), dtype=roots_batch.dtype)
    coeffs[:, -1] = 1
    for i in range(n_samples):
        for r in roots_batch[i]:
            for j in range(n_roots - 1, -1, -1):
                coeffs[i, j] = coeffs[i, j] - r * coeffs[i, j + 1]
    return coeffs

# 调用方式
coeffs = batch_polyfromroots(roots)

科学计算Python代码优化通用经验 #

• 优先替换Python层显式循环为numpy/pandas的向量化运算，向量化运算在C层执行，可大幅降低解释器开销。
• 对于无法向量化的计算逻辑，优先用Numba做JIT编译，改动量极小就能获得接近原生C的性能。
• 非必要场景下降低数据精度，比如用float32替代float64，既减少内存占用也能提升计算速度。
• 重复计算的公共逻辑提前提取，避免每轮循环都重复计算相同内容。
• 性能要求极高的场景可以用Cython编写静态类型的扩展模块，性能上限更高但开发成本也更高。

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