Selecting from 5D numpy array with a corresponding 3D array containing indices of the 4th dimension-CodePudding

I have a 5D numpy array containing values, and would like to obtain a subarray with one less dimension, where the values have been selected based on a 3D array that contains indices of the forth dimension of the first array. E.g., I have the following arrays:

values = np.random.randn(3,4,5,10,2)
indices = np.random.randint(0,values.shape[3],size=values.shape[:3])

I found one solution, but find it rather complicated, and would prefer a one-liner:

x = np.arange(values.shape[0])
y = np.arange(values.shape[1])
z = np.arange(values.shape[2])

result = values[x[:,None,None],y[None,:,None],z[None, None,:],indices,:]

Is there any better solution to get this array?

CodePudding user response：

You can try the following:

indices = indices[..., None, None]
result = np.take_along_axis(values, indices, axis=3).squeeze(axis=3)