I have an array a of size (M, N, K).
And I have an array b of size (M, N) with integer values of [0, K-1].
How do I get the array... c of size (M, N), where
c[i, j] == a[i, j, b[i, j]]
in the simplest manner?
Which part of the indexing guide is it?
You can use advanced indexing:
c = a[np.arange(M)[:, None], np.arange(N), b]
Output:
array([[ 0, 6, 12, 18],
[24, 25, 31, 37],
[43, 49, 50, 56]])
Taking @Vitalizzare's example:
# 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'):
a[*np.meshgrid(range(M), range(N), indexing='ij'), b]
See Advanced Indexing in NumPy and numpy.meshgrid for details.
Example:
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))
a[np.array(...), np.array(...), np.array(...)]. But thought it's a difficult way, too. Not far from c = np.array([ a[i, j, b[i,j]] for i in range(M) for j in range(N) ]).reshape(M, N)? But without loops and thanks about the reminder!
a[*np.indices(b.shape).reshape(2, -1), b.ravel()].reshape(*b.shape)c = np.array([ a[i, j, b[i,j]] for i in range(M) for j in range(N) ]).reshape(M, N). Post as a separate answer?