How to get partial cumulative sums (of positive and negative numbers) in an array?

Question

I have an array with positive and negative numbers and want to do a cumulative sum of numbers of the same sign until the next number carries an opposite sign. It starts again at 0. Maybe better explained with a sample.
Here is the original array:

np.array([0.2, 0.5, 1.3, 0.6, -0.3, -1.1, 0.2, -2.0, 0.7, 1.1, 0.0, -1.2])

And the output I expect without using a loop, of course:

np.array([0.0, 0.0, 0.0, 2.6, 0.0, -1.4, 0.2, -2.0, 0.0, 0.0, 1.8, -1.2])

Any efficient idea would help a lot…

Asked By: tibibou

||

Source

Answer 1

One vectorial option:

a = np.array([0.2, 0.5, 1.3, 0.6, -0.3, -1.1, 0.2, -2.0, 0.7, 1.1, 0.0, -1.2])

cs = np.cumsum(a)
idx = np.nonzero(np.r_[np.diff(a>0), True])
out = np.zeros_like(a)

out[idx] = np.diff(np.r_[0, cs[idx]])

Output:

array([ 0. ,  0. ,  0. ,  2.6,  0. , -1.4,  0.2, -2. ,  0. ,  1.8,  0. ,  -1.2])

Answered By: mozway

How to get partial cumulative sums (of positive and negative numbers) in an array?

Question:

Answers: