-2

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...

tibibou
  • 164
  • 10

1 Answers1

2

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])
mozway
  • 194,879
  • 13
  • 39
  • 75
  • Thank you mozway, it works as expected. How would you handle the last element in the a array (-1.2) so that it appears in the out array? – tibibou Nov 27 '22 at 14:35
  • Thank you mozway, it works with my original data and with 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, -1.6]). Unfortunately, it does not work if the last 2 numbers have a different sign as in : 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, 1.6]) I can add a 0 at the end, after 1.6, it works. Later on I can remove it in the out, but I suspect there is a smarter approach... – tibibou Nov 27 '22 at 15:15
  • @tibibou can you check again? – mozway Nov 27 '22 at 21:10