如何使用NumPy向量化实现无for循环的条件累加求和操作

你需要实现的是遇到0值就重置累加器的分段前缀和，不需要写Python层for循环，完全可以通过NumPy原生向量化操作实现，大数组下性能远高于显式循环。

可直接运行的代码 #

import numpy as np

# 原始测试数组
a = np.array([0, 1, 3, 0, 0, 5, 12, 1, 0, 6])

zero_mask = a == 0
# 第一步：将零值替换为0后计算全局前缀和
cumsum = np.where(zero_mask, 0, a).cumsum()
# 第二步：生成分组偏移量，开头补0兼容数组首元素非零的场景
offsets = np.r_[0, cumsum[zero_mask]]
# 第三步：给每个被零值分隔的非零段分配独立分组ID
groups = zero_mask.cumsum()
# 第四步：全局前缀和减去对应分组的偏移量，得到分段累加结果
b = cumsum - offsets[groups]

运行后得到的b输出为：

array([ 0,  1,  4,  0,  0,  5, 17, 18,  0,  6])

和你给出的目标结果完全一致，该写法同时兼容数组开头、结尾为非零值的场景。

逻辑拆解 #

1. 先把所有零值位置标记出来，将零值替换为0后计算全局累加值，此时零值位置的累加值和前一个非零段的最终累加值相等
2. 提取所有零值位置的累加值作为分组偏移量，在最前面补0作为第一个非零段的偏移基准
3. 通过零值掩码的累加值给每个分段分配独立ID，同一个非零段内的ID一致
4. 每个位置的全局累加值减去所属分段的偏移量，就得到了段内从0开始累加的结果，零值位置计算结果自动为0，完全匹配你需要的计算逻辑

所有操作都是NumPy底层C实现的向量化运算，没有Python层面的循环开销。

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