根据索引矩阵对NumPy矩阵进行排序

Question

在NumPy中，我们可以通过以下方式对数组进行排序：

>>> import numpy as np

>>> a = np.array([0, 100, 200])
>>> order = np.array([1, 2, 0])
>>> print(a[order])
[100 200   0]

然而，当“order”是一个矩阵时，这种方法就行不通了：

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

>>> Ord = np.array([  [1, 0, 2],
                      [0, 2, 1],
                      [2, 1, 0]])

>>> print(A[Ord].shape)
(3, 3, 3)

我希望将"A"按照以下方式排序：

array([[1, 0, 2],
       [3, 5, 4],
       [8, 7, 6]])

- Felip

Answer 1

你可以使用 np.take_along_axis 来实现此功能。

np.take_along_axis(A, Ord, axis=1)

输出

array([[1, 0, 2],
       [3, 5, 4],
       [8, 7, 6]])

如文档所述，它通常与生成索引的函数一起使用，例如argsort。但我不确定是否适用于超过2个维度。

- Kevin

Answer 2

有两种常见的方法可以索引到 numpy 数组：

1. 基本索引和切片，它扩展了 Python 的切片概念。
2. 高级索引，我们通常使用另一个包含整数/布尔值的 ndarray。

您可以在 numpy 文档 中了解这些内容。

回答您的问题：我们可以在这里使用高级索引。文档中包含一个示例（已经很接近我们想要的），如下所示：

>>> import numpy as np
>>> A = np.array([[1,2],
                  [3,4]])
>>> row_indices = np.array([[0,0],
                            [1,1]])
>>> col_indices = np.array([[1,0],
                            [1,0]])
>>> A[row_indices,col_indices]
array([[2, 1],
       [4, 3]])

在问题的代码中，Ord 已经包含了列索引，所以我们只需要自己生成行索引即可。虽然这可能不是最好的方法，但这里有一个可能的解决方案：

>>> A = np.array( [[0,1,2],
...                [3,4,5],
...                [6,7,8]])
>>> col_indices = np.array([[1,0,2],
...                         [0,2,1],
...                         [2,1,0]])
>>> row_indices = np.repeat([0,1,2],3).reshape(3,3)
>>> row_indices
array([[0, 0, 0],
       [1, 1, 1],
       [2, 2, 2]])
>>> A[row_indices, col_indices]
array([[1, 0, 2],
       [3, 5, 4],
       [8, 7, 6]])

- hasref