Numpy assignment of 3D array values using 2D index array

I have a 3D numpy array of zeros, with dimensions CxHxW (in this example, C=4, H=2, and W=3):

A = np.array([[[0, 0, 0],
               [0, 0, 0]],
              [[0, 0, 0],
               [0, 0, 0]]
              [[0, 0, 0],
               [0, 0, 0]]
              [[0, 0, 0],
               [0, 0, 0]]])

I also have a 2D array of indices, with dimensions HxW, such that every value in the array is a valid index between [0, C-1]

B = np.array([[2, 3, 0], 
              [3, 1, 2]])

Is there a fast way, using vectorization, to modify array A such that A[B[i][j]][i][j] = 1, for all valid i, j?

A = np.array([[[0, 0, 1],
               [0, 0, 0]],
              [[0, 0, 0],
               [0, 1, 0]]
              [[1, 0, 0],
               [0, 0, 1]]
              [[0, 1, 0],
               [1, 0, 0]]]) 

Thank you!

Answer

It seems like you are looking for put_along_axis:

np.put_along_axis(A, B[None,...], 1, 0)

Note that the second argument is required to have the same number of dimensions as the first, which is why B[None,...] is used instead of B.