基于布尔数组获取numpy数组的总和

Question

I have two numpy arrays prods and index

prods = np.asarray([ 0.5 ,  0.25,  1.98,  2.4 ,  2.1 ,  0.6 ])
index = np.asarray([False,  True,  True, False, False,  True], dtype=bool)

I need to calculate the sum of the values in prods array using the index array. The output I want to is

res = [0.75, 1.98, 5.1]

The first True in index array is preceded by a False , so I take the first two elements from prods (.5,.25) and sum them up(0.75). The second True in index has no preceding False (since its preceded by a True , the False at position zero doesn't count), so I simply output 1.98 in this case. The third True is preceded by two False , so I take those values from prods array (2.4,2.1,0.6) and sum them up. Any ideas on how to do this?

I basically need something like np.cumsum but I need to return the cumulative sum every time a True occurs in index and reset the cumulative sum value to zero.

Answer 1

You could use np.split and using np.where of your index array as positions to split:

>>> [arr.sum() for arr in np.split(prods, np.where(index)[0]+1)[:-1]]
[0.75, 1.98, 5.0999999999999996]

The last one isn't exactly 5.1 because of floating point precision. If you don't want to use Fraction s or Decimal s there's nothing you can do about that.

---

You could also use np.add.reduceat here:

>>> np.add.reduceat(prods, np.append([0], (np.where(index)[0]+1)[:-1]))
array([ 0.75,  1.98,  5.1 ])