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],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 flattened L similarly, indexed the first coordinate by L and reshaped back. This seems kind of stupid -- I am doing it like this because I don't really understand numpy indexing.
a, b, c = 5, 7, 4
A = np.random.rand(a, b, c)
# L indexes the first axis of A,
# so pick random indices that are lower than `a`:
L = np.random.randint(a, size=(a, b))
# Solution: create some indexers for the other axes:
j, k = np.indices((b, c), sparse=True)
# The broadcasting works in this case, because `L` indexes the first axis.
# In other situations, the following might be more robust:
# _, j, k = np.indices((a, b, c), sparse=True)
# Finally do the indexing:
B = A[L[:, :, None], j, k]
Checking:
B2 = np.zeros_like(A)
for i in range(a):
for j in range(b):
for k in range(c):
B2[i, j, k] = A[L[i, j], j, k]
assert np.array_equal(B, B2)