Selecting entries from slices with Numpy

Question

Suppose that we have array A of shape (3, 10, 10) and array B of shape (10, 10). Each element of B is index of an element along first axis of A. How can I create array C using operations where C[i, j] = A[B[i,j], i, j]?

Answer 1

Use advanced indexing and broadcasting:

import numpy as np

L, M, N = 3, 8, 10
A = np.arange(L)[:, None, None] + 10*np.arange(M)[:, None] + 100*np.arange(N)
B = np.random.randint(0, L, (M, N))
m, n = np.ogrid[:M, :N]
C = A[B, m, n]
np.all(C%10 == B)
# True