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

Question

3

我有两个多维数组，我想使用其中一个作为索引来生成一个新的多维数组。例如：

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

这可能很简单，但我在搜索中找不到答案。谢谢。

- cracka31

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Jaime · 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)。