有没有一种方法可以使用numpy向量化条件求和？

Question

我拥有的数组：

a = np.array([0, 1, 3, 0, 0, 5, 12, 1, 0, 6])

我需要的数组：

b = np.array([0, 1, 4, 0, 0, 5, 17, 18, 0, 6])

使用for循环创建数组b

b = np.zeros(a.size)
b[0] = a[0]

for i in range(1, a.size):
    if a[i] > 0:
        b[i] = a[i] + b[i-1]
    else:
        b[i] = 0

有没有一种方式可以向量化这种操作，而不需要使用for循环？

- Philippus Mercator

---

你的例子 "a" 和 "b" 在语法上不可行。 - Thomas Hilger

在你的例子中，a[i] < 0永远不会为真，对吗？ - Thomas Hilger

抱歉，我会纠正代码。 - Philippus Mercator

@ThomasHilger 我已经调整了代码，对此我很抱歉。 - Philippus Mercator

提供的实现存在一个错误，如果第一个元素>0，则应为b[0] = a[0]。 numba.jit是一个微不足道的选项，因为代码已经准备好编译（对于较大的数组，速度提高了约300倍）。或者使用@numba.vectorize和def gt_zero(a，b)：return a + b if b > 0 else 0进行自定义ufunc，使用gt_zero.accumulate(a)（约175倍加速）。需要numba ~0.55.1，因为在numba 0.51.1中引入了一个错误。即使可能，具有内在递归函数的向量化numpy也很难以合理的努力实现。 - Michael Szczesny

Answer 1

虽然不是严格的向量化，但可能仍有用。

import numpy as np

a = np.array([0, 1, 3, 0, 0, 5, 12, 1, 0, 6])

b = np.r_[*[np.cumsum(c) for c in np.split(a, np.where(a==0)[0])]]

print(b)

- Thomas Hilger