How can I extract elements from a 3D array using a 2D index array in NumPy?

You can use advanced indexing:

PYTHONCopied
c = a[np.arange(M)[:, None], np.arange(N), b]

Output:

PYTHONCopied
array([[ 0,  6, 12, 18],
       [24, 25, 31, 37],
       [43, 49, 50, 56]])

Taking @Vitalizzare's example:

PYTHONCopied
# input
M, N, K = 3, 4, 5

a = np.arange(M*N*K).reshape(M, N, K)
b = np.arange(M*N).reshape(M, N) % K

# output
array([[ 0,  6, 12, 18],
       [24, 25, 31, 37],
       [43, 49, 50, 56]])

To extract data this way you pass corresponding indexes as 2-dimentional arrays in each of 3 dimensions, where the first 2 should be the result of meshgrid for matrices (note indexign='ij'):

PYTHONCopied
   a[*np.meshgrid(range(M), range(N), indexing='ij'), b]

See Advanced Indexing in NumPy and numpy.meshgrid for details.

Example:

PYTHONCopied
import numpy as np

M, N, K = 3, 4, 5
a = np.arange(M*N*K).reshape(M, N, K)
b = np.arange(M*N).reshape(M, N) % K
first, second = np.meshgrid(range(M), range(N), indexing='ij')
c = a[first, second, b]

assert all(a[i, j, b[i,j]] == c[i,j] for i in range(M) for j in range(N))