python - 根据索引数组对值求和

Question

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

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

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

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

所以结果应该是 [10 + 3, 50 + 51 + 2, 6] 的 3 维数组

是否有一个 numpy 的方法来做到这一点？我有一些非常大的数组，我需要在多个维度上像这样求和，因此必须执行显式循环并不方便。

Answer 1

numpy.bincount有一个 weights 参数，它可以满足您的需要:

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