python - 有没有办法用 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 循环的情况下对这种类型的操作进行矢量化？

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)