Philippus Mercator
Philippus Mercator

Reputation: 43

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

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

Answers (1)

Thomas Hilger
Thomas Hilger

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

Related Questions