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

Question

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

You can use advanced indexing: arr[np.arange(len(arr)), :, indices]: numpy.org/doc/stable/reference/… — akuiper
– akuiper, Commented Dec 20, 2021 at 19:41
Thanks @Psidom - this works. do you want to add this as an answer? happy to accept it — jkr
– jkr, Commented Dec 20, 2021 at 19:44

ryangooch · Accepted Answer · 2021-12-20 19:45:18Z

3

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

answered Dec 20, 2021 at 19:45

ryangooch

543 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

user17242583 · Accepted Answer · 2021-12-20 19:49:02Z

1

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

answered Dec 20, 2021 at 19:49

user17242583

Collectives™ on Stack Overflow

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

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related