如何在Numpy中低内存高效计算向量与范德蒙德矩阵的乘积

你要实现的np.dot(A, np.vander(B, increasing=True))本质是计算加权幂和数组：结果的第k个元素为 $\sum_{i=0}^{N-1} A[i] \cdot B[i]^k$（k从0到N-1）。原方法因构造N×N范德蒙德矩阵导致O(N²)内存占用，大N场景下会触发内存不足，以下是两种O(N)内存的高效实现：

方法一：递推式计算（简单直接，O(N)内存） #

利用幂次的递推关系，仅维护当前B的k次幂数组（长度N）和结果数组（长度N），总内存复杂度为O(N)。Python层面的循环仅控制递推次数，核心运算依赖numpy优化的向量乘法：

import numpy as np

def compute_weighted_powers(A, B):
    N = len(A)
    res = np.empty(N, dtype=A.dtype)
    current_power = np.ones_like(B)
    res[0] = A @ current_power  # 计算k=0的幂和
    for k in range(1, N):
        current_power *= B       # 递推得到B的k次幂
        res[k] = A @ current_power
    return res

该方法内存占用极低，对于1e5级别的N也能轻松处理，且numpy内部的向量运算已做过高度优化，实际运行效率并不低。

方法二：无显式循环的JIT编译实现 #

如果严格要求避免Python层面循环，可使用numba对递推逻辑进行JIT编译，既保留O(N)内存效率，又能获得接近C语言的运行速度：

from numba import jit
import numpy as np

@jit(nopython=True)
def compute_weighted_powers_jit(A, B):
    N = len(A)
    res = np.empty(N, dtype=A.dtype)
    current_power = np.ones_like(B)
    res[0] = np.dot(A, current_power)
    for k in range(1, N):
        current_power *= B
        res[k] = np.dot(A, current_power)
    return res

注：若尝试纯numpy无循环实现，需借助初等对称多项式与线性递推，但会引入额外的复杂度，且对于大N场景，递推式方法的综合效率更优。

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