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

Question

我有两个 numpy 数组 prods 和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)

我需要使用index 数组计算prods 数组中值的总和。我想要的输出是

res = [0.75, 1.98, 5.1]

index 数组中的第一个True 前面有一个False，所以我从prods(.5,.25) 中取出前两个元素并将它们相加(0.75)。索引中的第二个True 没有前面的False（因为它前面有True，所以零位置的False 不计算在内），所以在这种情况下我只输出1.98。第三个True 前面有两个False，所以我从prods 数组(2.4,2.1,0.6) 中获取这些值并将它们相加。关于如何做到这一点的任何想法？

我基本上需要np.cumsum 之类的东西，但每次在索引中出现True 时，我都需要返回累积和，并将累积和值重置为零。

Answer 1

您可以使用np.split 并使用np.where 数组中的np.where 作为要拆分的位置：

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

由于浮点精度，最后一个不完全是5.1。如果您不想使用Fractions 或Decimals，那么您无能为力。

---

你也可以在这里使用np.add.reduceat：

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