使用索引列表对numpy数组的元素执行操作

Question

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?

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

Answer 1

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

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

Answer 2

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