Search code examples
numpymatrix-indexing

Indexing a 4D array using another array of 3D indices

A have a 4D array M (a x b x c x d) and an array I of indices (3 x f), e.g.

I = np.array([1,2,3, ...], [2,1,3, ...], [4,1,6, ...])

I would like to use I to arrive at a matrix X that has f rows and d columns, where:

X[0,:] = M[1,2,4,:]
X[1,:] = M[2,1,1,:]
X[2,:] = M[3,3,6,:]
...

I know I can use M[I[0], I[1], I[2]], however, I was wondering if there's a more concise solution?

Solution

  • You can use use, for example:

    I = np.array([[1,2,3], [2,1,3], [4,1,6]])
M = np.ndarray((10,10,10,10))
X = np.array([M[t,:] for t in I])