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

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

score 2 · Accepted Answer

这有点棘手，但以下将做到这一点：

>>> 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]会清楚地返回你所追求的。您可以像上面那样依靠广播获得相同的结果，将rows数组替换为np.arange(3)[:, np.newaxis]，相当于np.arange(3).reshape(3, 1).

1 回答 1