How can I use a 3d array of indices for a 2d array slicing in Numpy

Question

I have 2 arrays as input. On array as output. Array a holds the data and is of shape (N,M) , while array b holds the indices and is of shape (N,X,2) . The resulting array should be of shape (N,X) , with the values taken from a .

Right now it only works with a for loop. How could I vectorize it since I have huge arrays as input?

Below is a sample code to demonstrate what I have right now:

import numpy as np

# a of shape (N,M)
# b of shape (N,X,2)
# t_result of shape (N, X)

a = np.random.randint(0, 10, size=(10, 10))
b = np.random.randint(0, 2, size=(10, 9, 2))

t_result = np.empty((10, 9))

for i in range(b.shape[0]):
    t_result[i] = a[i, b[i, :, 0]]

print(t_result)
print(t_result.shape)

Answer 1

Ok so I adapted a bit the answer to another post from scleronomic :

import numpy as np

# a of shape (N,M)
# b of shape (N,X,2)
# t_result of shape (N, X)

a = np.random.randint(0, 10, size=(10, 10))
b = np.random.randint(0, 2, size=(10, 9, 2))

t_result = np.empty((10, 9))

t_result = a[np.arange(a.shape[0])[:,None],b[np.arange(b.shape[0]),:,0]]

print(t_result)
print(t_result.shape)

I am not sure whether or not it is the best solution but it works.