Reputation: 43
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?
Upvotes: 3
Views: 137
Reputation: 476
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)
Upvotes: 1