numpy,按1d索引数组选择行中的元素

Question

我们有方阵,n*n.例如,n = 3,数组是这样的:

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

让我们在每个ROW中都有一系列索引.例如:

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

我想得到:

[1,5,7]

因为在行[0,1,2]中取索引为1的元素,在行[3,4,5]中获取索引为2的元素,在行[6,7,8]中获取索引为1的元素.

我很困惑,并且不能使用标准的numpy索引以这种方式获取元素.谢谢你的答案.

Answer 1

没有真正漂亮的方式但这会做你想要的:)

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