Increase numpy array elements using array as index

Question

I am trying to efficiently update some elements of a numpy array A, using another array b to indicate the indexes of the elements of A to be updated. However b can contain duplicates which are ignored whereas I would like to be taken into account. I would like to avoid for looping b. To illustrate it:

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

whereas I would like the output to be:

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

Any ideas?

Answer 1

To correctly handle the duplicate indices, you'll need to use np.add.at instead of +=. Therefore to update the first row of A, the simplest way would probably be to do the following:

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

The documents for the ufunc.at method can be found here.