获取 NumPy 数组中多个索引处的元素

Question

正如documentation 解释的那样，我正在尝试访问 NumPy 数组的特定行和列，但我认为我遗漏了一些东西。

我有以下数组：

arr = np.random.randint(10, size=(6, 4))
array([[1, 9, 6, 4],
       [8, 5, 0, 3],
       [3, 7, 3, 2],
       [1, 4, 8, 0],
       [5, 5, 8, 0],
       [0, 6, 4, 9]])

我想得到第一行和最后一行；以及第一、第三和最后一列，所以我正在尝试：

arr[(0, -1),(0, 1, 3)]

但这会产生以下错误：

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (2,) (3,)

我想我误解了这种整数索引。我希望这个输出：

array([[1, 9, 4],
       [0, 6, 9]])

我可以这样做但是感觉很尴尬：

arr[(0,-1),:][:,(0,1,3)]

如何获取不同维度的第 i 个元素？

Answer 1

一种方法是显式指定结果的每个元素的索引，其中关键区别在于行规范应该是 2d：

>>> np.random.seed(444)
>>> arr = np.random.randint(10, size=(6, 4))
>>> arr
array([[3, 0, 7, 8],
       [3, 4, 7, 6],
       [8, 9, 2, 2],
       [2, 0, 3, 8],
       [0, 6, 6, 0],
       [3, 0, 6, 7]])

>>> rows = [[0],
...         [-1]]
>>> cols = [0, 1, -1]
>>> arr[rows, cols]
array([[3, 0, 8],
       [3, 0, 7]])

此示例利用行索引的广播。换句话说，您可以利用行在除第 0 列之外的其余两列中重复自身这一事实。

要使用高级索引，需要明确选择所有元素。 ...但是，由于上面的索引数组只是重复自己，因此可以使用广播。您无法利用广播的示例，需要指定每个单独的索引，例如：

rows = [[1, 2],
        [3, 4]]
cols = [[2, 3],
        [1, 0]]

在这种情况下，您选择的是 (1, 2), (2, 3), ... 处的元素

请参阅文档中的 this 示例。

您可能会注意到这正是np.ix_ 产生的内容（尽管我并不是说它不是一个有用的功能）：

>>> rows, cols = np.ix_((0, -1),(0, 1, 3))

Answer 2

您正在寻找np.ix_:

arr[np.ix_((0, -1),(0, 1, 3))]

结果

array([[1, 9, 4],
       [0, 6, 9]])