使用数组作为索引增加numpy数组元素

Question

我试图有效地更新numpy数组A的一些元素,使用另一个数组b来指示要更新的A元素的索引.但是b可以包含被忽略的重复项,而我希望将其考虑在内.我想避免循环b.为了说明它:

>>> A = np.arange(10).reshape(2,5)
>>> A[0, np.array([1,1,1,2])] += 1
>>> A
array([[0, 2, 3, 3, 4],
       [5, 6, 7, 8, 9]])

而我希望输出为:

array([[0, 3, 3, 3, 4],
       [5, 6, 7, 8, 9]])

有任何想法吗？

Answer 1

要正确处理重复索引,您需要使用np.add.at而不是+=.因此,要更新第一行A,最简单的方法可能是执行以下操作:

>>> np.add.at(A[0], [1,1,1,2], 1)
>>> A
array([[0, 4, 3, 3, 4],
       [5, 6, 7, 8, 9]])

该ufunc.at方法的文件可以在这里找到.