numpy，按一维索引数组选择行中的元素

Question

假设我们有一个正方形的数组，大小是 n*n。比如说 n=3，数组长这样：

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

然后，我们在每一行都有一个索引数组。比如说：

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

更简单的方法是使用choose这个numpy函数：

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

Answer 2

没有特别好看的方法，不过这个可以满足你的需求 :)

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