Numpy: Indexing 3D matrix using 1D array

Question

arr = np.arange(12).reshape((3, 2, 2))
indices = np.array([0, 1, 1])

expected_outcome = np.array([[0, 1], [6, 7], [10, 11]])

I'm trying to index this array of shape (3,2,2) with an array of shape (3) containing the y-index of the value I want to get. I tried to make it work with for in statement, but is there an elegant way to do it with numpy?

I'd probably do np.array([arr[i][indices[i]] for i in range(arr.shape[0])]), but like you said, it would be nice to vectorize that, and it seems like something that would be doable. — ichthyophile
– ichthyophile, Commented Jul 25, 2022 at 18:53

hpaulj · Accepted Answer · 2022-07-25 19:25:47Z

1

So you want arr[0,0,:], arr[1,1,:], arr[2,1,:]?

How about

In [179]: arr[[0,1,2], [0,1,1]]
Out[179]: 
array([[ 0,  1],
       [ 6,  7],
       [10, 11]])

answered Jul 25, 2022 at 19:25

hpaulj

233k14 gold badges260 silver badges392 bronze badges

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

2 Comments

ichthyophile Over a year ago

Do you know of a way to use the indices variable without a for-loop?

hpaulj Over a year ago

@ichthyophile,do you wmatn arr[np.arange(3), indices]. The key is that the index for a nd arrray should be a tuple of arrays - arrays that broadcast against each other.

Collectives™ on Stack Overflow

Numpy: Indexing 3D matrix using 1D array

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related