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!
Advertisement
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
.