CODEWITHSUNDEEP
pythonarraysnumpy
geompalik
geompalik

Reputation: 1582

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?

Upvotes: 6

Views: 3161

Answers (2)

Alex Riley
Alex Riley

Reputation: 176770

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.

Upvotes: 11

TheBlackCat
TheBlackCat

Reputation: 10298

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]

Upvotes: 0

Related Questions