使用从np.argsort（）返回的索引对2D numpy数组进行排序

Question

当我们具有一维numpy数组时，可以按以下方式对其进行排序：

>>> temp = np.random.randint(1,10, 10)
>>> temp
array([5, 1, 1, 9, 5, 2, 8, 7, 3, 9])
>>> sort_inds = np.argsort(temp)
>>> sort_inds
array([1, 2, 5, 8, 0, 4, 7, 6, 3, 9], dtype=int64)
>>> temp[sort_inds]
array([1, 1, 2, 3, 5, 5, 7, 8, 9, 9])

注意：我知道我可以使用np.sort; 显然，我需要其他数组的排序索引-这只是一个简单的示例。现在我们可以继续我的实际问题。

我试图对2D数组应用相同的方法：

>>> d = np.random.randint(1,10,(5,10))
>>> d
array([[1, 6, 8, 4, 4, 4, 4, 4, 4, 8],
       [3, 6, 1, 4, 5, 5, 2, 1, 8, 2],
       [1, 2, 6, 9, 8, 6, 9, 2, 5, 8],
       [8, 5, 1, 6, 6, 2, 4, 3, 7, 1],
       [5, 1, 4, 4, 4, 2, 5, 9, 7, 9]])
>>> sort_inds = np.argsort(d)
>>> sort_inds
array([[0, 3, 4, 5, 6, 7, 8, 1, 2, 9],
       [2, 7, 6, 9, 0, 3, 4, 5, 1, 8],
       [0, 1, 7, 8, 2, 5, 4, 9, 3, 6],
       [2, 9, 5, 7, 6, 1, 3, 4, 8, 0],
       [1, 5, 2, 3, 4, 0, 6, 8, 7, 9]], dtype=int64)

这个结果看起来不错-请注意，我们可以d使用sort_inds1D示例中演示的相应行的索引对每一行进行排序。但是，尝试使用我在1D示例中使用的相同方法来获得排序后的数组，但出现了以下异常：

>>> d[sort_inds]
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-63-e480a9fb309c> in <module>
----> 1 d[ind]

IndexError: index 5 is out of bounds for axis 0 with size 5

所以我有两个问题：

1. 刚刚发生了什么？numpy如何解释此代码？
2. 我如何仍能使用实现我想要的-即排序d-或任何其他相同维度的数组sort_inds？

谢谢

Answer 1

您需要一些额外的工作才能正确索引二维数组。这是使用高级索引的方法，其中np.arange在第一个轴中使用 where ，以便sort_indsin 中的每一行从 中 的相应行中提取值d：

d[np.arange(d.shape[0])[:,None], sort_inds]

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