Search code examples
pythonnumpymultidimensional-arraymatrix-indexing

How to do a simple ndarray indexing operation

I'm trying to use numpy's "indexing with an array" feature.

Given this code:

import numpy as np
x = np.array([[10, 20, 30], [40, 50, 60]])
index = ???
print(x[index])

I want to print: [10, 50]

That is, I'd like to find the correct shape array that will fetch from two successive rows in a matrix.

index = np.array([0, 1])
index = np.array([[0], [1]])

Both print this: [10, 20] and [[10], [20]]

For these cases, all of the index rows are being "broadcasted" to the first row of the target. This makes no sense. What is the correct technique to index out of successive rows of the matrix?

Solution

  • You can pass both two dimensions 

    x[np.arange(len(x)),[0,1]]
Out[137]: array([10, 50])