numpy, Select elements in rows by 1d indexes array

Question

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.

I am pretty sure that there is a better way, but for a start: numpy.array([arr[i, j] for (i, j) in enumerate(myidx)]) — Nras
– Nras, Commented Sep 9, 2014 at 6:57

Wolph · Accepted Answer · 2014-09-09 07:06:48Z

9

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 Sep 9, 2014 at 7:06

Wolph

80.4k12 gold badges142 silver badges152 bronze badges

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

2 Comments

dondublon Over a year ago

Cool! This really pretty, don't doubt!

Wolph Over a year ago

@QtRoS: numpy slicing supports multiple dimensions so the arange() generates the x indices and the myidx contains the y indices. So it's slicing like this: arr[x_idx, y_idx]

belgraviton · Accepted Answer · 2021-06-07 14:22:32Z

1

Simpler way to reach the goal is using choose numpy function:

numpy.choose(myidx, arr.transpose())

answered Jun 7, 2021 at 14:22

belgraviton

1541 silver badge3 bronze badges

Collectives™ on Stack Overflow

numpy, Select elements in rows by 1d indexes array

2 Answers 2

2 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related