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

时间:2019-05-03 08:54:39

标签: python arrays numpy multidimensional-array

当我们拥有一维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_inds的每一行进行排序,如一维示例所示。但是,尝试使用与一维示例相同的方法来获取排序数组,但出现了以下异常:

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

所以我有2个问题:

  1. 发生了什么事? numpy如何解释此代码?
  2. 如何仍能使用d实现我想要的-即对sort_inds-或任何其他相同尺寸的数组进行排序?

谢谢

1 个答案:

答案 0 :(得分:4)

您需要做一些额外的工作才能正确索引2d数组。这是使用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]])