numpy，按一維索引數組選擇行中的元素

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

Answer 2

達到目標的更簡單方法是使用選擇numpy 函數：

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