在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]

Answer 2

原则上，您可以使用numpy的{​​{1}}和bincount进行此操作，但我猜它只会使代码的可读性降低，而不会有任何明显的性能提升。