Index numpy 3d-array with 1d array of indices-CodePudding

I have a 3D numpy array of shape (i, j, k). I have an array of length i which contains indices in k. I would like to index the array to get a shape (i, j).

Here is an example of what I am trying to achieve:

import numpy as np
arr = np.arange(2 * 3 * 4).reshape(2, 3, 4)

# array([[[ 0,  1,  2,  3],
#         [ 4,  5,  6,  7],
#         [ 8,  9, 10, 11]],
#
#       [[12, 13, 14, 15],
#        [16, 17, 18, 19],
#        [20, 21, 22, 23]]])

indices = np.array([1, 3])

# I want to mask `arr` using `indices`
# Desired output is equivalent to
# np.stack((arr[0, :, 1], arr[1, :, 3]))
# array([[ 1,  5,  9],
#        [15, 19, 23]])

I tried reshaping the indices array to be able to broadcast with arr but this raises an IndexError.

arr[indices[np.newaxis, np.newaxis, :]]
# IndexError: index 3 is out of bounds for axis 0 with size 2

I also tried creating a 3D mask and applying it to arr. This seems closer to the correct answer to me but I still end up with an IndexError.

mask = np.stack((np.arange(arr.shape[0]), indices), axis=1)
arr[mask.reshape(2, 1, 2)]
# IndexError: index 3 is out of bounds for axis 0 with size 2

CodePudding user response：

From what I understand in your example, you can simply pass indices as your second dimension slice, and a range of length corresponding to your indices for the zeroth dimension slice, like this:

import numpy as np
arr = np.arange(2 * 3 * 4).reshape(2, 3, 4)

indices = np.array([1, 3])

print(arr[range(len(indices)), :, indices])
# array([[ 1,  5,  9],
#        [15, 19, 23]])

CodePudding user response：

This works:

sub = arr[[0,1], :, [1,3]]

Output:

>>> sub
array([[ 1,  5,  9],
       [15, 19, 23]])

A more dynamic version by @Psidom:

>>> sub = arr[np.arange(len(arr)), :, [1,3]]
array([[ 1,  5,  9],
       [15, 19, 23]])