两次argsort就能直接拿到样本位次？靠谱吗，性能咋样？

数组排序

在处理各类数据过程中，对数组进行排序是一种非常常见的操作。在实际应用场景中，我们会用到sort()、argsort()等多种排序函数。一般而言，sort()函数返回的是，一个按升序规则，对原数组中元素按数值大小进行排序后的数组；而argsort()的返回值则是sort()函数返回的数组中的元素，在原数组中的实际位置。在这里，我以 Python 中的 Numpy 包为例，演示这两个函数的用法

1
2
3

num = [0,3,2,1]
print(np.sort(num))
print(np.argsort(num))

结果为：

1
2

[0 1 2 3]
[0 3 2 1]

获取排名

除了上述两种排序场景外，还有另一种稍稍复杂一些的排序场景。假设在一个项目中，有 10000 个样本需要进行处理和评估，评估结果是一个长度为 10000 的数组。而对于这个评估结果，我不仅需要将最后结果排序，还需要知道每个样本的评估结果，在所有评估结果中的具体排名位次是多少。假如我们依旧是使用 Python 与 Numpy 对数据进行处理，那我们最直观的一个做法就是，对数据使用argsort()函数排序后，再对结果进行遍历，将排名依次对应到具体的样本上，具体代码如下所示：

1
2
3
4
5
6
7

num = [7,3,5,2,4,8,6]
ast = np.argsort(num)
rank = np.zeros_like(num)
for index in range(len(ast)):
    rank[ast[index]] = index
print(ast)
print(rank)

计算结果为：

1
2

[3 1 4 2 6 0 5]
[5 1 3 0 2 6 4]

这种计算方式从思路上来讲是最为直观的，计算复杂度也不高，排序部分为$O(n\log{n})$，获取排名部分的复杂度为$O(n)$。

复杂度越低就一定越好吗？

一般来讲，一个$O(n)$的算法，从执行效率上来讲，是要高于一个$O(n\log{n})$的算法的，除非在一些特殊场景下。

例如在小数据量下的 HashMap 与 TreeMap，尽管 HashMap 的查找复杂度为$O(1)$而 TreeMap 为$O(\log{n})$，但在数据量较小时，对数据进行 Hash 运算反而比树查找要慢。

在使用 Python 进行数据处理时其实也是这样，在使用 Numpy 进行运算时，由于 Numpy 底层运算完全由 C 编写，同时还使用了 BLAS 库进行加速，在进行矩阵与向量运算时效率极高，远远超出原生 Python 的性能。而 Python 部分由于 Python 是一个包装的奇奇怪怪的脚本语言，其本质还是基于一个逐行执行的解释器来运行的，缺乏编译优化的能力，再加上 GC 对性能的影响，在运行时的直观影响就是算法复杂度为$O(n\log{n})$的排序部分算起来没准比复杂度为$O(n)$的获取排名部分慢不了多少，甚至还要更快一些。

例如当我们对一个随机生成的长度为 10000 的数据进行排位时，统计各部分的耗时结果：

1
2
3
4
5
6
7
8
9

num = np.random.randn(100000)
start = time.time()
ast = np.argsort(num)
print("排序耗时", time.time() - start)
start = time.time()
rank = np.zeros_like(num)
for index in range(len(ast)):
    rank[ast[index]] = index
print("获取排名耗时", time.time() - start)

结果如下：

1
2

排序耗时 0.0069997310638427734
获取排名耗时 0.023000717163085938

结果复杂度为$O(n)$的获取排名部分耗时反而比复杂度为$O(n\log{n})$的排序部分高了 2 倍多。

优雅的排位实现

上面介绍的那种先排序，再根据排序结果统计位次的方法，虽然思路上很直观，但这个代码实现实在太过于丑陋了。加之面临原生 Python 的性能问题，因此我们可以采用另一种更加优雅，在实际运行中速度也不见得会更慢的实现手法。

这种实现方法其实非常简单，对原始的评估结果嵌套应用两层argsort函数就可以得到每个样本的评估结果在所有评估结果中的排名位次了，实例代码如下：

1
2
3

num = [7,3,5,2,4,8,6]
rank = np.argsort(np.argsort(num))
print(rank)

运行结果为

1

[5 1 3 0 2 6 4]

这个实现看起来有些违背直觉，为啥嵌套两层argsort就能直接拿到排名位次呢，在下一小节将给出对这一结论的证明

证明过程

以下证明针对升序排序过程进行，降序排序同理也可推导得出

对于一维数组num，假设有数组ast=argsort(num)与2ast=argsort(ast)，而对于任意数组a，a.idx(i)表示在数组a中，值为i的元素的 index。

基于这些设置，我们想要证明的内容可以写作：

$$ \forall{i,j\in \mathrm{2ast}}, i>j \Rightarrow \mathrm{num}[\mathrm{2ast}.\mathbf{idx}(i)] > \mathrm{num}[\mathrm{2ast}.\mathbf{idx}(j)] $$

在正式证明之前，关于argsort函数，我们可以给出引理如下：

$$ \begin{aligned} \textbf{Lemma:}\quad & 对于任意数组\mathrm{a}满足\mathbf{set}(\mathrm{a})==\mathbf{set}(\mathbf{range}(0,n))且\mathbf{len}(\mathrm{a})==n。\\\\
& \mathbf{argsort}(\mathrm{a})[i]为\mathrm{a}中第i大的元素的index \\
& 又由于\mathrm{a}满足的性质，可知\mathrm{a}[\mathbf{argsort}(\mathrm{a})[i]]==i一定成立 \end{aligned} $$

证明过程如下

$$ \begin{aligned} \begin{array}{r l} \textbf{Proof:}\quad & \forall{i,j\in \mathrm{2ast}}, i>j \Longrightarrow \mathrm{2ast}[\mathrm{2ast}.\mathbf{idx}(i)]>\mathrm{2ast}[\mathrm{2ast}.\mathbf{idx}(j)]\\\\
& 由引理可知， \mathrm{ast}[\mathrm{2ast}[i]]=i\Longrightarrow \mathrm{ast}.\mathbf{idx}(i)=\mathrm{2ast}[i] \\\\
& \mathrm{ast}.\mathbf{idx}(i)=\mathrm{2ast}[i]\Longrightarrow \mathrm{ast}.\mathbf{idx}(\mathrm{2ast}.\mathbf{idx}(i))=\mathrm{2ast}[\mathrm{2ast}.\mathbf{idx}(i)] \\\\
& \therefore \forall{i,j\in \mathrm{2ast}}, i>j \Longrightarrow \mathrm{ast}.\mathbf{idx}(\mathrm{2ast}.\mathbf{idx}(i)) > \mathrm{ast}.\mathbf{idx}(\mathrm{2ast}.\mathbf{idx}(j))\\\\
又& \because 根据\mathbf{argsort}函数性质，\forall {i,j\in \mathrm{ast}}, \mathrm{ast}.\mathbf{idx}(i) > \mathrm{ast}.\mathbf{idx}(j) \Longrightarrow \mathrm{num}[i]>\mathrm{num}[j]\\\\
且& \forall {i\in \mathrm{2ast}}, 始终有\mathrm{2ast}.\mathbf{idx}(i)\in \mathrm{ast}\\\\
& \therefore \forall {i,j\in \mathrm{2ast}}, i>j\Rightarrow \mathrm{num}[\mathrm{2ast}.\mathbf{idx}(i)] > \mathrm{num}[\mathrm{2ast}.\mathbf{idx}(j)]\\
\end{array} \end{aligned} $$

根据如上过程，可以证明，嵌套两层argsort就能直接拿到样本的排名位次。

遍历算法的另一种实现

除了使用原生 Python 以外，还可以使用 Numpy 来实现遍历排位。这种方法需要使用 Numpy 的高级索引功能，具体实现如下：

1
2
3
4
5
6

num = [7,3,5,2,4,8,6]
ast = np.argsort(num)
rank = np.zeros_like(num)
rank[ast] = np.arange(len(num))
print(ast)
print(rank)

计算结果为：

1
2

[3 1 4 2 6 0 5]
[5 1 3 0 2 6 4]

在这种方法中，第四行中使用 argsort 到的结果，将位次 0 到 n 依次赋给对应的排位 index 值。这一实现看起来没有双层 argsort 简洁，但速度上更有优势。

效率咋样？

相比单层argsort后遍历结果，双层argsort在获取排名部分的算法复杂度为$O(n\log n)$，相比直接遍历的$O(n)$复杂度，从算法的角度来看其实并不高效。但由于 Numpy 对底层算法的高效实现与封装，在 Python 上实现这一功能时，当数组的长度不是非常大时，双层argsort的方式都是更快的。同时，利用这种方法，在 TensorFlow 或 PyTorch 等计算框架下，同样可以获得较好的通用性与加速。

为了比较不同数据规模下，不同方法的运算效率，我随机生成了规模从 10 到 10 亿的数据模拟测试各种方法的性能，具体结果如下表所示：

结果表明，在数据规模不超过 10 亿时，双层argsort的方法都是要优于使用原生 Python 遍历排位的方法的。但如果将遍历排位部分更改为 Numpy 实现，则在大多数情况时能获得比双层argsort稍高的速度。而在数量不超过 10000 时，由于 GPU 与内存之间的通信延迟等问题，使用 TensorFlow 加速双层argsort方法并没有太大优势，但当数据规模继续扩大时，TensorFlow 的加速效果便十分明显了。