python - 在 numpy 中向量化索引操作

Question

在 numpy 中向量化以下代码的最佳方法是什么？

from numpy import *

A = zeros(5, dtype='int')
I = [1, 1, 1, 3]
J = [2, 1, 1, 1]
for i, j in zip(I, J):
    A[i] += j

print A

结果应该是:

[0 4 0 1 0]

在这里A是原始数组，I存储我们要增加 J 对应条目的索引.

如果通过以下方式简单地将上述内容矢量化:

A[I] += J
print A

答错了

[0 1 0 1 0]

显然，重复的索引被忽略了。是否有与 += 等效的操作？哪个不忽略重复索引？

Answer 1

你可以使用numpy.bincount():

A = numpy.zeros(5, dtype='int')
I = [1, 1, 1, 3]
J = [2, 1, 1, 1]
sums = numpy.bincount(I, J)
A[:len(sums)] += sums
print(A)

打印

[0 4 0 1 0]