arrays - numpy 在另一个多维数组上使用多维索引数组

Question

我有一个 2 个多维数组，我想用一个作为索引来生成一个新的多维数组。例如:

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

b = array([[0,2],[1,1],[3,1]])

我想用b中的第一个数组返回a中第一个数组中的那些索引元素，依此类推。所以我希望输出是:

array([[4,2],[8,8],[6,1]])

这可能很简单，但我无法通过搜索找到答案。谢谢。

Answer 1

这有点棘手，但以下内容可以做到:

>>> a[np.arange(3)[:, np.newaxis], b]
array([[4, 2],
       [8, 8],
       [6, 1]])

您需要索引 a 数组的行和列，因此要匹配您的 b 数组，您需要这样的数组:

rows = np.array([[0, 0],
                 [1, 1],
                 [2, 2]])

然后 a[rows, b] 会清楚地返回您想要的内容。依赖广播可以得到和上面一样的结果，用np.arange(3)[:, np.newaxis]替换rows数组，相当于np.arange(3).reshape(3, 1).