NumPy 中的排序方法（sort, argsort, lexsort, partition, argpartition, searchsorted）

---

sort

numpy.sort(arr, axis = -1, kind = None, order = None)

参数的具体含义我们稍后都会提到。

和列表类型一样，NumPy arrays 也能使用 sort 方法在原地进行排序：

import numpy as np

arr = np.random.randn(6)
arr
"""
array([-1.39195481,  0.73836815,  1.96817565, -0.69126633, -0.83675666,
       -0.32110115])
"""
1
2
3
4
5
6
7
8

arr.sort()
arr
"""
array([-1.39195481, -0.83675666, -0.69126633, -0.32110115,  0.73836815,
        1.96817565])
"""
1
2
3
4
5
6

axis 参数和其它 NumPy 函数的含义是一样的，用来指定沿哪个维度进行排序。默认值为 -1，表示沿最后一个轴的方向进行排序。

arr = np.random.randn(5, 3)
arr
"""
array([[-0.71873827, -1.79714254,  0.07859657],
       [-0.42673412,  0.82879723, -0.54362253],
       [ 0.79449046,  0.08115767, -0.34500705],
       [ 0.4613742 ,  0.15215783,  0.10526428],
       [-0.71593027,  0.55641392, -0.70285329]])
"""
1
2
3
4
5
6
7
8
9

arr.sort(1)
arr
"""
array([[-1.79714254, -0.71873827,  0.07859657],
       [-0.54362253, -0.42673412,  0.82879723],
       [-0.34500705,  0.08115767,  0.79449046],
       [ 0.10526428,  0.15215783,  0.4613742 ],
       [-0.71593027, -0.70285329,  0.55641392]])
"""
1
2
3
4
5
6
7
8
9

如果我们要使用上层函数 mp.sort，那么会返回原始 array 排好序的拷贝，而不是在原地进行排序。

arr = np.random.randn(1000)
quantile_5 = np.sort(arr)
quantile_5[int(0.05 * len(quantile_5))] # %5 quantile
"""
-1.7157833501117195
"""
1
2
3
4
5
6

sort 方法对原始数组升序，并没有直接的参数选项来帮助我们进行降序。我们可以使用和列表 [::-1]一样的 trick 来达到降序的目的：

arr = np.random.randn(3, 5)
arr.sort(axis=1)
arr
"""
array([[-2.16450452, -0.60159095, -0.20855381,  0.13868267,  0.94555186],
       [-0.96156166, -0.91536779, -0.14838184,  0.66383755,  1.11318848],
       [-1.59858316, -1.26248709, -0.31035371,  0.10956498,  2.20005666]])
"""
1
2
3
4
5
6
7
8

arr[:, ::-1]
"""
array([[ 0.94555186,  0.13868267, -0.20855381, -0.60159095, -2.16450452],
       [ 1.11318848,  0.66383755, -0.14838184, -0.91536779, -0.96156166],
       [ 2.20005666,  0.10956498, -0.31035371, -1.26248709, -1.59858316]])
"""
1
2
3
4
5
6

---

argsort, lexsort

接下来我们看一下 kind 参数，我们平时使用很少会考虑这个参数，但是还是很有必要了解一下的。kind 参数指定排序时使用的算法，我们列出几个可用的算法及其性能，熟悉数据结构与算法的朋友一定不会感到陌生：

稳定（stable）的排序算法说明，排序前后相同元素的相对顺序不会发生改变。这在我们直接使用 sort 的情景中很少考虑。但在考虑间接排序（indirect sort）时，元素的相对顺序可能很重要。

例如，我们有如下待排序的值：

values = np.array(['2:first', '2:second', '1:first', '1:second', '1:third'])
1

我们根据如下 key 值来排序：

key = np.array([2, 2, 1, 1, 1])
1

我们可以通过归并排序这种稳定的排序算法来排序：

indexer = key.argsort(kind='mergesort')
indexer
"""
array([2, 3, 4, 0, 1], dtype=int64)
"""
1
2
3
4
5

argsort 返回给我们排好序之后的元素在原来 array 的对应位置下标。

values.take(indexer)
"""
array(['1:first', '1:second', '1:third', '2:first', '2:second'],
      dtype='
1
2
3
4
5

对于一组学生姓名的数据，我们可能想先根据姓来排序，然后再根据名字来排序，这就是间接排序。例如：

first_name = np.array(['Bob', 'Jane', 'Steve', 'Bill', 'Barbara'])
last_name = np.array(['Jones', 'Arnold', 'Arnold', 'Jones', 'Walters'])

sorter = np.lexsort((first_name, last_name))
sorter
"""
array([1, 2, 3, 0, 4], dtype=int64)
"""
1
2
3
4
5
6
7
8

[last_name[i] + " " + first_name[i] for i in sorter]
"""
['Arnold Jane', 'Arnold Steve', 'Jones Bill', 'Jones Bob', 'Walters Barbara']
"""
1
2
3
4

虽然 last_name 似乎是后传入 lexsort 中的，但 lexsort 会先对 last_name 进行排序，这点要注意。lexsort 排序是稳定的。

---

partition, argpartition

我们排序最重要的目的可能就是为了找到最大值或者最小值。partition 方法会创建传入 array 的拷贝，且重置数据位置，使得最小的 k 个元素在最前面。至于这 k 个元素的顺序以及其它元素的顺序则并未定义。

arr = np.random.randn(20)
arr
"""
array([-2.04970854,  0.55915432, -0.53417685,  0.65999909, -0.56440136,
       -1.81955786,  0.38976505,  0.72029978, -0.68116799, -0.69901343,
       -0.05395895, -0.12968292, -0.34663156, -1.69320134,  0.25935788,
        0.83778976, -0.64741487,  0.08517585,  0.38104945,  0.37615883])
"""
1
2
3
4
5
6
7
8

np.partition(arr, 5)
"""
array([-1.81955786, -2.04970854, -1.69320134, -0.69901343, -0.68116799,
       -0.64741487, -0.56440136, -0.53417685, -0.34663156, -0.12968292,
       -0.05395895,  0.65999909,  0.55915432,  0.72029978,  0.25935788,
        0.83778976,  0.38976505,  0.08517585,  0.38104945,  0.37615883])
"""
1
2
3
4
5
6
7

和 argsort 类似，argpartition 返回的是索引值：

indices = np.argpartition(arr, 3)
indices
"""
array([ 5,  0, 13,  9,  4,  3,  6,  7,  8,  1, 10, 11, 12,  2, 14, 15, 16,
       17, 18, 19], dtype=int64)
"""
1
2
3
4
5
6

arr.take(indices)
"""
array([-1.81955786, -2.04970854, -1.69320134, -0.69901343, -0.56440136,
        0.65999909,  0.38976505,  0.72029978, -0.68116799,  0.55915432,
       -0.05395895, -0.12968292, -0.34663156, -0.53417685,  0.25935788,
        0.83778976, -0.64741487,  0.08517585,  0.38104945,  0.37615883])
"""
1
2
3
4
5
6
7

---

searchsorted

searchsorted 在一个有序数组 a 上通过二分搜索来为我们寻找传入元素 v 应该插入的位置：

numpy.searchsorted(a, v, side='left', sorter=None)

arr = np.array([0, 1, 7, 12, 15])
arr.searchsorted(9)
"""
3
"""
1
2
3
4
5

我们也可以传入一个 array，会返回 array 中每个元素应该插入的位置：

arr.searchsorted([0, 8, 11, 16])
"""
array([0, 3, 3, 5], dtype=int64)
"""
1
2
3
4

注意 0 应该插入的位置为 0，这是因为 side 参数的设置，默认为 'left'，表明插入元素应满足：

$$a[i-1]<v\le a[i]$$
通俗点的说，如果 v 在 a 中已有相同元素，那么 side='left' 会选择这些相同元素最左边的位置插入，side='right' 同理：

arr = np.array([0, 0, 0, 1, 1, 1, 1])
arr.searchsorted([0, 1])
"""
array([0, 3], dtype=int64)
"""
1
2
3
4
5

arr.searchsorted([0, 1], side='right')
"""
array([3, 7], dtype=int64)
"""
1
2
3
4

再来看 searchsorted 的一个应用：

data = np.floor(np.random.uniform(0, 10000, size=50))
1

我们想对 data 按如下区间分类：

bins = np.array([0, 100, 500, 1000, 5000, 10000])
1

使用 searchsorted 我们可以轻松得到 data 中每个值所属的区间（1 表示 [0, 100)）:

labels = bins.searchsorted(data)
labels
"""
array([5, 4, 5, 4, 4, 5, 2, 5, 5, 4, 5, 5, 2, 5, 4, 4, 4, 4, 4, 5, 4, 4,
       5, 3, 5, 4, 4, 4, 4, 5, 5, 5, 1, 5, 4, 5, 3, 4, 4, 5, 5, 5, 4, 5,
       5, 4, 3, 4, 2, 4], dtype=int64)
"""
1
2
3
4
5
6
7

---

References

[1] NumPy Reference. https://numpy.org/doc/stable/reference/index.html
[2] Python for Data Analysis, 2${}^{\mathrm{n}\mathrm{d}}$ edition. Wes McKinney.