按索引计算的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))的数组,即使某些计数为零

Answer 2

此解决方案对于大型数组应该更有效（它迭代可能的索引值而不是 的各个条目i）：

import numpy as np

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

i_max = i.max()
c = np.empty(i_max+1)
for j in range(i_max+1):
    c[j] = d[i==j].sum()

print c
[1. 2. 7.]