Perform operation on elements of numpy array using indexes list

I have numpy array and two python lists of indexes with positions to increase arrays elements by one. Do numpy has some methods to vectorize this operation without use of for loops?

for

My current slow implementation:

a = np.zeros([4,5])
xs = [1,1,1,3]
ys = [2,2,3,0]

for x,y in zip(xs,ys): # how to do it in numpy way (efficiently)?
a[x,y] += 1

print(a)


Output:

[[0. 0. 0. 0. 0.]
[0. 0. 2. 1. 0.]
[0. 0. 0. 0. 0.]
[1. 0. 0. 0. 0.]]


2 Answers
2

np.add.at will do just that, just pass both indexes as a single 2D array/list:

np.add.at

a = np.zeros([4,5])
xs = [1, 1, 1, 3]
ys = [2, 2, 3, 0]

np.add.at(a, [xs, ys], 1) # in-place
print(a)

array([[0., 0., 0., 0., 0.],
[0., 0., 2., 1., 0.],
[0., 0., 0., 0., 0.],
[1., 0., 0., 0., 0.]])


>>> a = np.zeros([4,5])
>>> xs = [1, 1, 1, 3]
>>> ys = [2, 2, 3, 0]
>>> a[[xs,ys]] += 1
>>> a
array([[ 0., 0., 0., 0., 0.],
[ 0., 0., 1., 1., 0.],
[ 0., 0., 0., 0., 0.],
[ 1., 0., 0., 0., 0.]])


By clicking "Post Your Answer", you acknowledge that you have read our updated terms of service, privacy policy and cookie policy, and that your continued use of the website is subject to these policies.