如何用Numpy对数组迭代应用无法通过ufunc构建的自定义函数？

好问题！确实Numpy没有像TensorFlow的tf.scan那样原生支持带状态的累积迭代函数，但我们有几个比纯Python循环更高效、更符合Numpy风格的解决方案：

方案1：生成器 + numpy.fromiter #

利用Python生成器处理状态累积，再通过numpy.fromiter直接生成Numpy数组——这个方法避免了手动逐个赋值数组元素，fromiter会在C层面完成数组创建，比纯循环快很多：

import numpy as np

r_list = np.arange(1,6,dtype=np.float32)
gamma = 0.99

def pv_accumulator(sequence, gamma, initial_state=0.0):
    current_state = initial_state
    for val in sequence:
        current_state = val + current_state * gamma
        yield current_state

# 从生成器直接生成Numpy数组
R_list = np.fromiter(pv_accumulator(r_list, gamma), dtype=np.float32)
print(R_list)
# 输出结果：[ 1.        2.99      5.9601    9.900499  14.801494]

方案2：用Numba加速循环（推荐大规模数组） #

如果你可以引入第三方库，Numba能把Python循环编译成机器码，速度接近Numpy原生函数的水平，完美适配这种递归式的状态累积场景：

import numpy as np
from numba import jit

r_list = np.arange(1,6,dtype=np.float32)
gamma = 0.99

# 用Numba装饰器编译函数
@jit(nopython=True)
def compute_pv(r_array, gamma_val):
    result = np.zeros_like(r_array)
    state = 0.0
    for i in range(len(r_array)):
        state = r_array[i] + state * gamma_val
        result[i] = state
    return result

R_list = compute_pv(r_list, gamma)
print(R_list)
# 输出结果和之前一致

为什么这比原生循环好？ #

你原来的循环是逐个修改Numpy数组元素，这会带来不必要的Python层开销。而fromiter是一次性完成数组初始化，Numba则直接把循环逻辑编译成高效的机器码，两者在数组规模越大时，性能优势越明显。

Numpy的核心优势是向量化运算，但对于这种依赖前一步状态的递归操作，确实没有专门的ufunc支持，上面的方案就是当前最适合的Numpy友好型实现了。

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