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的元素。 / p>

我很困惑，并且不能使用标准的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

实现目标的更简单方法是使用 choose numpy 函数：

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

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

2 个答案: