Python Numpy无循环向量化实现依赖前序值的递推数组计算

你的代码问题分析 #

你写的重复赋值逻辑无效，是因为numpy向量化切片赋值会先计算完所有右值再统一写入左值，不会逐行更新第二列的结果给后续行的计算使用，只有显式循环才会逐行更新左值，所以这种写法无法实现递推逻辑。

实现原理 #

我们可以通过数学变换把递推式转换为可向量化的累加操作：
给定递推规则：

b[i] = 0.5 * b[i-1] + 0.5 * a[i] （i≥1，b[0]为初始值，a[i]为第i行第一列的值）
两边同乘2^i得到：
2^i * b[i] = 2^{i-1} * b[i-1] + 2^{i-1} * a[i]
令c[i] = 2^i * b[i]，则递推式简化为c[i] = c[i-1] + 2^{i-1}*a[i]，c[0] = b[0]，这是一个纯累加计算，可以直接用numpy的cumsum函数向量化实现。

完整实现代码 #

import numpy as np

# 初始数组
a = np.asarray([[1,2,3],[4,np.nan,np.nan],[7, np.nan, np.nan], [8, np.nan, np.nan]])
n = a.shape[0]

# 生成2的幂次数组
pow2 = 2 ** np.arange(n)
# 构造c数组的初始值
c = np.zeros(n)
c[0] = a[0, 1]
c[1:] = pow2[:-1] * a[1:, 0]
# 累加得到完整c数组
c = np.cumsum(c)
# 转换回第二列的结果并赋值
a[:, 1] = c / pow2

# 输出结果
print(a)

运行输出 #

[[1.  2.  3. ]
 [4.  3.  nan]
 [7.  5.  nan]
 [8.  6.5 nan]]

该方案完全没有Python级别的循环，所有计算都是Numpy底层的向量化操作，即使数组长度很大也能保持很高的执行效率。

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