How to vectorize increments in Python

Question

I have a 2d array, and I have some numbers to add to some cells. I want to vectorize the operation in order to save time. The problem is when I need to add several numbers to the same cell. In this case, the vectorized code only adds the last. 'a' is my array, 'x' and 'y' are the coordinates of the cells I want to increment, and 'z' contains the numbers I want to add.

import numpy as np

a=np.zeros((4,4))
x=[1,2,1]
y=[0,1,0]
z=[2,3,1]
a[x,y]+=z
print(a)

As you see, a[1,0] should be incremented twice: one by 2, one by 1. So the expected array should be:

[[0. 0. 0. 0.]
 [3. 0. 0. 0.]
 [0. 3. 0. 0.]
 [0. 0. 0. 0.]]

but instead I get:

[[0. 0. 0. 0.]
 [1. 0. 0. 0.]
 [0. 3. 0. 0.]
 [0. 0. 0. 0.]]

The problem would be easy to solve with a for loop, but I wonder if I can correctly vectorize this operation.

Answer 1

Use np.add.at for that:

import numpy as np

a = np.zeros((4,4))
x = [1, 2, 1]
y = [0, 1, 0]
z = [2, 3, 1]
np.add.at(a, (x, y), z)
print(a)
# [[0. 0. 0. 0.]
#  [3. 0. 0. 0.]
#  [0. 3. 0. 0.]
#  [0. 0. 0. 0.]]