[FIXED] numpy, Select elements in rows by 1d indexes array

Issue

Let we have square array, n*n. For example, n=3 and array is this:

arr = array([[0, 1, 2],
   [3, 4, 5],
   [6, 7, 8]])

And let we have array of indices in the each ROW. For example:

myidx=array([1, 2, 1], dtype=int64)

I want to get:

[1, 5, 7]

Because in line [0,1,2] take element with index 1, in line [3,4,5] get element with index 2, in line [6,7,8] get element with index 1.

I'm confused, and can't take elements this way using standard numpy indexing. Thank you for answer.

Solution

There's no real pretty way but this does what you are looking for :)

In [1]: from numpy import *

In [2]: arr = array([[0, 1, 2],
   [3, 4, 5],
   [6, 7, 8]])

In [3]: myidx = array([1, 2, 1], dtype=int64)

In [4]: arr[arange(len(myidx)), myidx]
Out[4]: array([1, 5, 7])

Answered By - Wolph