Numpy: Using an index array to set values in a 3D array

Question

I have an indices array of shape (2, 2, 3) which looks like this:

array([[[ 0,  6, 12],
        [ 0,  6, 12]],
       [[ 1,  7, 13],
        [ 1,  7, 13]]])

I want to use these as indices to set some values of a np.zeros matrix to 1. While the highest value in this example is a 13, I know that it can go up to 18. This is why I created one_hot = np.zeros((2, 2, 18)) array:

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, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
       [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]])

Using the indices array, my desired outcome is this:

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

I want to use numpy's advanced indexing sort of like this:

one_hot[indices] = 1

How can I do that?

Answer 1

A possible solution:

np.put_along_axis(one_hot, indices, 1, axis=2)

Output:

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

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