根据索引数组求和值

Question

我有两个相同维度的数组：

a = np.array([ 1,  1, 2,  0, 0, 1])
b = np.array([50, 51, 6, 10, 3, 2])

我想根据b中的索引对a的元素求和。

我想要的矩阵的i元素将是所有值b[j]的总和，使a[j]==i。

因此结果应该是[10 + 3, 50 + 51 + 2, 6]

的三维数组

有这么简单的方法吗？我有一些非常大的数组，我需要在多个维度上这样求和，所以不得不执行显式循环。

Answer 1

numpy.bincount有一个weights参数，可以满足您的需求：

In [36]: np.bincount(a, weights=b)
Out[36]: array([  13.,  103.,    6.])

Answer 2

如果你没有使用numpy，那么简单如下：

res = [0]*len(set(a))
for i, v in enumerate(b):
    res[a[i]] += v

假设a中的索引总是从0开始并且是连续序列。