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]])