How to do indexing of a NumPy 3D-array based on 2D-array in Python?

Question

Let say I have a NumPy array A of shape (66,5) and B of shape (100, 66, 5).

The elements of A will index the first dimension (axis=0) of B, where the values are from 0 to 99 (i.e. the first dimension of B is 100).

A = 
array([[   1,    0,    0,    1,    0],
       [   0,    2,    0,    2,    4],
       [   1,    7,    0,    5,    5],
       [   2,    1,    0,    1,    7],
       [   0,    7,    0,    1,    4],
       [   0,    0,    3,    6,    0]
       ....                         ]])

For example, A[4,1] will take index 7 of the first dimension of B, index 4 of the second dimension of B and index 1 of the third dimension B.

What I wanted to is to produce array C of shape (66,5) where it contains the elements in B that are selected based on the elements in A.

Answer 1

You can use np.take_along_axis to do that:

import numpy as np
np.random.seed(0)
a = np.random.randint(100, size=(66, 5))
b = np.random.random(size=(100, 66, 5))
c = np.take_along_axis(b, a[np.newaxis], axis=0)[0]
# Test some element
print(c[25, 3] == b[a[25, 3], 25, 3])
# True