python - 通过重复数组索引对数组值求和

Question

我想将 vals 中的值求和到索引列表 idx 中指定的较小数组 a 的元素中。

import numpy as np

a = np.zeros((1,3))
vals = np.array([1,2,3,4])
idx = np.array([0,1,2,2])

a[0,idx] += vals

这会产生结果 [[ 1. 2. 4.]] 但我想要结果 [[ 1. 2. 7.]]，因为它应该将 vals 中的 3 和 vals 中的 4 添加到 a 的第二个元素中。

我可以通过以下方式实现我想要的:

import numpy as np

a = np.zeros((1,3))
vals = np.array([1,2,3,4])
idx = np.array([0,1,2,2])

for i in np.unique(idx):
    fidx = (idx==i).astype(int)
    psum = (vals * fidx).sum()
    a[0,i] = psum 

print(a)

有没有办法在不使用 for 循环的情况下使用 numpy 执行此操作？

Answer 1

可以使用 np.add.at，只要形状对齐，即 a 在这里需要是一维的。

a = a.squeeze()
np.add.at(a, idx, vals)

a
array([1., 2., 7.])