通过索引累计求和numpy数组

Question

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

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

第二个数组指定哪些元素需要求和

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

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

我试着用

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

bincount返回所有整数范围（max（i））的数组，即使某些计数为零

Answer 2

朱的评论是最有效的解决方案。以下是工作代码：

import numpy as np
import scipy.ndimage as ni

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

n_indices = i.max() + 1
print ni.sum(d, i, np.arange(n_indices))

Answer 3

对于大型数组，此解决方案应该更有效（它迭代可能的索引值，而不是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.]