numpy如何无for循环实现依赖前序元素的递归向量化计算

以下分两类场景给出可直接落地的实现方案：

1. 纯NumPy原生方案（仅支持满足结合律的二元累积操作） #

如果你的自定义函数可以拆解为可结合的二元迭代操作（类似加法、乘法满足f(f(a,b),c) = f(a,f(b,c))的形式），可以直接用NumPy ufunc的accumulate()方法实现无Python层循环的累积计算，性能和内置的cumsum/cumprod持平。

示例代码 #

import numpy as np

# 示例：自定义累积逻辑为 前序结果 * 0.9 + 当前元素（指数加权平均）
def custom_op(prev, curr):
    return prev * 0.9 + curr

# 将自定义二元函数转为NumPy ufunc
custom_ufunc = np.frompyfunc(custom_op, nin=2, nout=1)

# 生成测试数组
Size = 10000
a = np.random.randn(Size)

# 累积计算，自动遍历所有元素，无需手动写循环
b = custom_ufunc.accumulate(a, dtype=np.object_).astype(np.float64)

注意：该方案仅支持可结合的二元操作，无法通用适配所有任意函数。

2. 通用高性能方案（无需编写C代码，适配任意自定义逻辑） #

如果你的函数逻辑复杂，无法拆解为可结合的二元操作，直接用Numba的JIT编译即可。你不需要修改原有for循环的逻辑，只需要加一个装饰器，Numba会自动将代码编译为机器码，性能和手写C扩展相当，完全规避Python层循环的开销。

示例代码 #

import numpy as np
from numba import njit

# 直接保留你原本的for循环逻辑，加@njit装饰器即可
@njit
def calc_custom_b(a):
    size = a.shape[0]
    b = np.empty(size, dtype=np.float64)
    # 自定义初始值，可根据需求调整
    b[0] = a[0]
    for i in range(1, size):
        # 此处可写入任意你需要的逻辑，包括调用np.exp、np.sum等NumPy函数
        # 示例：计算前i个元素的指数之和
        b[i] = np.sum(np.exp(a[:i]))
    return b

# 调用时直接传入数组即可
Size = 100000
a = np.random.randn(Size)
b = calc_custom_b(a)

Numba是Python数据科学生态的常用库，可直接通过pip install numba安装，无需额外配置C编译环境。

补充说明 #

不存在完全没有循环的实现方案，所有这类累积操作底层都需要遍历，我们要做的只是避免Python层的慢循环。上述两个方案都不需要你手动编写C/C++扩展，即可达到接近原生编译的性能。

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