python - 按索引对 numpy 数组进行累积求和

Question

假设您有一个需要求和的值数组

d = [1,1,1,1,1]

第二个数组指定哪些元素需要加在一起

i = [0,0,1,2,2]

结果将存储在一个大小为 max(i)+1 的新数组中。因此，例如 i=[0,0,0,0,0] 相当于对 d 的所有元素求和并将结果存储在位置 0 大小为 1 的新数组。

我尝试使用

c = zeros(max(i)+1)
c[i] += d

但是，+= 操作只添加每个元素一次，因此给出了意想不到的结果

[1,1,1]

代替

[2,1,2]

如何正确实现这种求和？

Answer 1

如果我没看错问题，有一个快速函数可以解决这个问题(只要数据数组是 1d)

>>> i = np.array([0,0,1,2,2])
>>> d = np.array([0,1,2,3,4])
>>> np.bincount(i, weights=d)
array([ 1.,  2.,  7.])

np.bincount 返回所有整数范围 (max(i)) 的数组，即使某些计数为零