如何在二维数组上使用 numpy.argsort 对另一个二维数组进行排序

Question

我一直使用 numpy.argsort 处理一维数据，但它在 2D 中的表现似乎有所不同。

例如，假设我想沿轴 1 对该数组进行 argsort，以便每行中的项目按升序排列

>>> import numpy as np
>>> arr = np.eye(4)
>>> arr
array([[1., 0., 0., 0.],
       [0., 1., 0., 0.],
       [0., 0., 1., 0.],
       [0., 0., 0., 1.]])

>>> idx = np.argsort(arr, axis=1)
>>> idx
array([[1, 2, 3, 0],
       [0, 2, 3, 1],
       [0, 1, 3, 2],
       [0, 1, 2, 3]])

到目前为止一切都很好。

上面的每一行给出了列在第二个数组中应该如何重新排列的顺序。

假设我们想用上面的idx对下面的数组进行排序。

>>> arr2 = np.arange(16).reshape((4, 4))
>>> arr2
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

>>> sorted = arr2[idx]
>>> sorted
array([[[ 8,  9, 10, 11],
        [12, 13, 14, 15],
        [ 0,  1,  2,  3],
        [ 4,  5,  6,  7]],
       ....
       [[ 8,  9, 10, 11],
        [12, 13, 14, 15],
        [ 0,  1,  2,  3],
        [ 4,  5,  6,  7]]])

>>> sorted.shape
(10, 4, 4)

形状现在增加了尺寸。

我期待得到。

array([[ 1,  2,  3, 0],
       [ 4,  6,  7, 5],
       [ 8,  9, 11, 10],
       [12, 13, 14, 15]])

我可以对行进行迭代，这很糟糕！

>>> rows = []
>>> for i, row in enumerate(arr2):
...     rows.append(row[idx[i]])
>>> np.arrays(rows)
array([[ 1,  2,  3,  0],
       [ 4,  6,  7,  5],
       [ 8,  9, 11, 10],
       [12, 13, 14, 15]])

Answer 1

np.take_along_axis 有一个使用argsort 的示例：

>>> a = np.array([[10, 30, 20], [60, 40, 50]])

We can sort either by using sort directly, or argsort and this function

>>> np.sort(a, axis=1)
array([[10, 20, 30],
       [40, 50, 60]])
>>> ai = np.argsort(a, axis=1); ai
array([[0, 2, 1],
       [1, 2, 0]])
>>> np.take_along_axis(a, ai, axis=1)
array([[10, 20, 30],
       [40, 50, 60]])

这简化了将ai 应用于数组本身的过程。我们可以直接这样做，但需要多考虑一下索引实际代表什么。

在此示例中，ai 是沿轴 1 的索引值（值如 0、1 或 2）。这个 (2,3) 必须 broadcast 与轴 0 的 (2,1) 数组：

In [247]: a[np.arange(2)[:,None], ai]
Out[247]: 
array([[10, 20, 30],
       [40, 50, 60]])