Adding numbers to 2-d array based on index

user_8253878 0
Viewed 24

I have a numpy array of size 21:

arr1
array([  0., 329., 730., 513.,   0., 167.,   0.,   0., 175.,   0., 220.,
         0.,   0.,   0., 202.,   0.,   0.,  59.,   0.,  33.,  47.])

I have an indexed array of size 21:

arg_arr
array([4, 3, 2, 3, 1, 3, 2, 0, 3, 0, 3, 2, 2, 1, 0, 4, 4, 3, 2, 0, 3],
      dtype=int64)

I need add the elements to a numpy array of zeros of size 5 based on their index. i.e. at index 0, the output arr2 = 0 + 0 + 202 + 33.

arr2 = np.zeros((5,))
array([0., 0., 0., 0., 0.])

How can I do this with numpy?

numpy arrays
1 Answers

This is a textbook usecase for np.add.at:

np.add.at(arr2, arg_arr, arr1)

The special thing about ufunc.at vs just doing arr2[arg_arr] += arr1 is that the operation is unbuffered, so multiple occurrences of an index are handled correctly.

user_29887300
Related