使用从 np.argsort() 返回的索引对 2D numpy 数组进行排序 [重复]

Question

当我们有一个 1D 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 做到这一点；显然，我需要不同数组的排序索引——这只是一个简单的例子。现在我们可以继续我的实际问题..

我尝试对二维数组应用相同的方法：

>>> 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)

这个结果看起来不错 - 请注意，我们可以使用来自sort_inds 的相应行的索引对d 的每一行进行排序，如一维示例所示。但是，尝试使用我在 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

您需要做一些额外的工作才能正确索引二维数组。这是使用advanced indexing的一种方式，其中np.arange用于第一个轴，以便sort_inds中的每一行从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]])