I have a numpy array A with shape (a, b, c), and another integer array L with shape (a, b). I want to make an array B such that B[i, j, k] = A[L[i, j], k] (assume shapes and values of L permit this).

What is the most efficient way to do this? The only ways I can think to do it involve creating meshgrids or new tiled arrays to expand L to the same shape as A which feels inefficient and not very nice. Surely this is quite a standard thing to do with an efficient implementation somewhere (is there an einops way?)?

I used einops einops.rearrange(A, i j k -> (i j) k)and flattenedLsimilarly, indexed the first coordinate byL` and reshaped back. This seems kind of stupid — I am doing it like this because I don’t really understand numpy indexing.