Increase numpy array elements using array as index

问题:

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?

回答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.

回答2:

One approach is to use numpy.histogram to find out how many values there are at each index, then add the result to A:

A[0, :] += np.histogram(np.array([1,1,1,2]), bins=np.arange(A.shape[1]+1))[0]