Correct indexing of a multidimensional Numpy array with another array of indices

Question

I am trying to index a multidimensional P array with another indices array. which determines which element along the last axis I want, as follows:

 import numpy as np M, N = 20, 10 P = np.random.rand(M,N,2,9) # index into the last dimension of P indices = np.random.randint(0,9,size=(M,N)) # I'm after an array of shape (20,10,2) # but this has shape (20, 10, 2, 20, 10) P[...,indices].shape 

How to index P with indices to get an array of form (20,10,2) ?

If this is not too clear: for any i and j (within the bounds) I want my_output[i,j,:] be equal to P[i,j,:,indices[i,j]]

Answer 1

I think this will work:

 P[np.arange(M)[:, None, None], np.arange(N)[:, None], np.arange(2), indices[..., None]] 

Not really, I know ...

---

This may look better, but it may also be less picky:

 P[np.ogrid[0:M, 0:N, 0:2]+[indices[..., None]]] 

or perhaps better:

 idx_tuple = tuple(np.ogrid[:M, :N, :2]) + (indices[..., None],) P[idx_tuple]