python - 添加到由另一个数组索引的数组的矢量化方式 - Python/NumPy

Question

我有 3 个 Numpy 数组，a、b 和 c。b 和c 是非常大的数组并且具有相同的长度。 b的每个元素都是0、1或2，a的长度也是3。现在我想知道是否有一种方法可以消除以下 for 循环:

for i in range(len(b)):
    a[b[i]] += c[i]

如有任何意见，我们将不胜感激。

Answer 1

NumPy ufunc 有一个 at这种情况的方法:

numpy.add.at(a, b, c)

这做了每个人都期望 a[b] += c 对索引数组 b 做的事情，然后才发现它不起作用。