Vectorizing an indexing operation in numpy

Question

What is the best way to vectorize the following code in 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

The result should be:

[0 4 0 1 0]

Here A is the original array, I stores the index at which we want to increment by the corresponding entry of J .

If one simply vectorizes the above by doing:

A[I] += J
print A

one gets the wrong answer

[0 1 0 1 0]

as, apparently, repeated indices are ignored. Is there an equivalent operation to += which does not ignore repeated indices?

Answer 1

You can use 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)

prints

[0 4 0 1 0]

Answer 2

原则上你可以用numpy的bincount和unique来做到这一点，但是我猜它只会使代码的可读性降低，而不会有任何合理的性能提升。