Is there a way to vectorize a conditional sum with numpy?

user_19427202 0
Viewed 83

array I have:

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

array I need:

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

for loop that gives me array 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

is there a way to vectorize this type of operation doing that without for loops?

vectorization conditional-statements numpy sum python
1 Answers

It is not exactly vectorized, but maybe useful nevertheless

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)
user_70845660
Related