Numpy之numpy.argsort()-CFANZ编程社区

Numpy之numpy.argsort()

前言

关于argsort这个函数一直没有很好的理解，今天算是搞明白了。特此记录。

介绍

函数原型：

numpy.argsort(a, axis=-1, kind=None, order=None)

理解这个函数的关键在于明白其返回的是什么。

首先明确，对于numpy的数组输出的形式上，横向始终表示的是最后一维的方向（axis=-1）。对于二维的数据a（np.ndim(a)=2），axis=0表示列的方向，也就是从上到下的方向；axis=1表示行的方向，也就是从左到右的方向。

它的返回值表示这么几点：

返回值中的元素表示原始数据的索引（也就是位置），直接返回排序后的结果可以使用numpy.sort()
返回值中的元素的位置顺序表示了真实数据的从小到大的排序顺序
返回值中的元素对应于元素值表示的索引上按照指定轴（方向）的原始数据

可以分析下面几个例子加深理解。

a
array([[0.3451194 , 0.32197821, 0.23369311],
       [0.57426324, 0.03603039, 0.46324327],
       [0.73053454, 0.25336478, 0.59913947],
       [0.47781611, 0.79507359, 0.11370495]])
# 沿着列排序，也就是沿着维度axis=0方向上排序
np.argsort(a, axis=0)
array([[0, 1, 3],
       [3, 2, 0],
       [1, 0, 1],
       [2, 3, 2]])
# 沿着行排序，也就是沿着维度axis=1方向上排序
np.argsort(a, axis=1)
array([[2, 1, 0],
       [1, 2, 0],
       [1, 2, 0],
       [2, 0, 1]])

可以看出来，对于numpy数组a而言，使用了np.argsort()后的返回结果，axis决定了排序针对的数据。

对于np.argsort(a, axis=0)的输出的[1, 0]位置的值3：

表示这个位置应该放置原始数据沿着axis=0方向上（列向）的坐标值为3的数据
也就是a[3, 0] = 0.4778161084937589
也即是[1, 0] => 3 => [3, 0] ==>> 排序后位置[1, 0]上的值为a[3, 0]

根据numpy文档的例子中可以看到，根据得到的索引是可以得到最终的排序结果的：

ind = np.argsort(a, axis=0)
ind
array([[0, 1, 3],
       [3, 2, 0],
       [1, 0, 1],
       [2, 3, 2]])
np.take_along_axis(a, ind, axis=0)
array([[0.3451194 , 0.03603039, 0.11370495],
       [0.47781611, 0.25336478, 0.23369311],
       [0.57426324, 0.32197821, 0.46324327],
       [0.73053454, 0.79507359, 0.59913947]])

这里使用np.take_along_axis()得到的结果和直接使用np.sort()效果是一样的。

np.sort(a, axis=0)
array([[0.3451194 , 0.03603039, 0.11370495],
       [0.47781611, 0.25336478, 0.23369311],
       [0.57426324, 0.32197821, 0.46324327],
       [0.73053454, 0.79507359, 0.59913947]])

参考资料

https://docs.scipy.org/doc/numpy/reference/generated/numpy.argsort.html