在 numpy 中使用列表和数组进行索引似乎不一致

Question

受this other question 的启发，我试图在 NumPy 中围绕advanced indexing 进行思考，并对其工作原理建立更直观的理解。

我发现了一个有趣的案例。这是一个数组：

>>> y = np.arange(10)
>>> y
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

如果我将它索引为标量，我当然会得到一个标量：

>>> y[4]
4

使用一维整数数组，我得到另一个一维数组：

>>> idx = [4, 3, 2, 1]
>>> y[idx]
array([4, 3, 2, 1])

所以如果我用一个二维整数数组索引它，我会得到......我会得到什么？

>>> idx = [[4, 3], [2, 1]]
>>> y[idx]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: too many indices for array

哦不！对称性被打破。我必须用 3D 数组索引才能得到 2D 数组！

>>> idx = [[[4, 3], [2, 1]]]
>>> y[idx]
array([[4, 3],
       [2, 1]])

是什么让 numpy 有这种行为？

为了让这更有趣，我注意到使用 numpy 数组（而不是列表）进行索引的行为符合我的直觉预期，而 2D 给了我 2D：

>>> idx = np.array([[4, 3], [2, 1]])
>>> y[idx]
array([[4, 3],
       [2, 1]])

这看起来与我所在的位置不一致。这里有什么规则？

Answer 1

原因是将列表解释为 numpy 数组的索引：列表被解释为元组，而使用元组的索引被 NumPy 解释为多维索引。

就像arr[1, 2]返回元素arr[1][2]一样，arr[[[4, 3], [2, 1]]]等同于arr[[4, 3], [2, 1]]，并且会根据多维索引的规则返回元素arr[4, 2]和arr[3, 1]。

通过再添加一个列表，您确实告诉 NumPy 您想要沿第一个维度进行切片，因为最外面的列表被有效地解释为好像您只传递了一个“第一个维度的索引列表”：arr[[[[4, 3], [2, 1]]]]。

来自documentation：

示例

应从每一行中选择一个特定元素。行索引只是 [0, 1, 2]，列索引指定要为相应行选择的元素，这里是 [0, 1, 0]。将两者结合使用，可以使用高级索引来解决任务：

>>> x = np.array([[1, 2], [3, 4], [5, 6]])
>>> x[[0, 1, 2], [0, 1, 0]]
array([1, 4, 5])

and:

警告

高级索引的定义意味着x[(1,2,3),] 与x[(1,2,3)] 根本不同。后者相当于x[1,2,3]，它将触发基本选择，而前者将触发高级索引。请务必了解发生这种情况的原因。

在这种情况下，最好使用np.take：

>>> y.take([[4, 3], [2, 1]])  # 2D array
array([[4, 3],
       [2, 1]])

这个函数 [np.take] 与“花式”索引（使用数组索引数组）做同样的事情；但是，如果您需要沿给定轴的元素，它会更容易使用。

或者将索引转换为数组。这样 NumPy 将其解释为（array 是特殊情况！）作为花哨的索引而不是“多维索引”：

>>> y[np.asarray([[4, 3], [2, 1]])]
array([[4, 3],
       [2, 1]])