Numpy get values at indices given in array form

Question

I want to get the values at indices of my_array.

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

my_array = np.array([[[1.1587323 , 1.75406635],
        [1.05464125, 1.29215026],
        [0.9784655 , 1.16957462]]])

I should get the following output:

output: array([[[1.1587323], [1.29215026], [0.9784655]]])

Is it possible without for loops or list comprehensions?

Answer 1

You can use np.take_along_axis:

np.take_along_axis(my_array, indices, axis=-1)
array([[[1.1587323 ],
        [1.29215026],
        [0.9784655 ]]])