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