Python 中 sorted 如何自定义比较逻辑

在 Python 中对一个可迭代对象进行排序是很常见的一个操作,一般会用到 sorted() 函数

num_list = [4, 2, 8, -9, 1, -3]
sorted_num_list = sorted(num_list)
print(sorted_num_list)

上面的代码是对整数列表 num_list 按从小到大的顺序进行排序,得到的结果如下

[-9, -3, 1, 2, 4, 8]

有时候不仅仅是对元素本身进行排序,而是在元素值的基础上进行一些计算之后再进行比较,比如将 num_list 中的元素按照其平方值的大小进行排序。

在 Python 2 中,可以通过 sorted() 函数中的 cmp 或 key 参数来实现这种自定义的比较逻辑。cmp 比较函数接收两个参数 x 和 y(x 和 y 都是列表中元素)并且返回一个数字,如果返回正数表示 x > y,返回 0 表示 x == y,返回负数表示 x < y。key 函数接收一个参数,重新计算出一个结果,然后用计算出的结果参与排序比较。因此在 Python 2 中按平方值大小排序可以有下面两种实现方式

num_list = [4, 2, 8, -9, 1, -3]
# cmp 参数只在 Python 2 中存在,Python 3 及之后的版本移除了 cmp 参数
sorted_num_list = sorted(num_list, cmp=lambda x, y: x ** 2 - y ** 2)
sorted_num_list = sorted(num_list, key=lambda x: x ** 2)

但是随着 Python 3.0 的发布,cmp 参数也随之被移除了,也就是说在 Python 3 中自定义比较逻辑就只能通过 key 参数来实现。至于为什么将 cmp 参数移除,在 Python 的 Issue tracker 中有一段很长的讨论,主要有以下两点原因

关于上面性能的问题,我做了一个实验,分别随机生成 1000、10000、100000 和 1000000 个整数,然后用 key 和 cmp 的方式分别进行排序并记录排序的时间消耗

import random
import time

counts = (1000, 10000, 100000, 1000000)

def custom_cmp(x, y):
    return x ** 2 - y ** 2

def custom_key(x):
    return x ** 2

print('%7s%20s%20s' % ('count', 'cmp_duration', 'key_duration'))
for count in counts:
    min_num = -count // 2
    max_num = count // 2
    nums = [random.randint(min_num, max_num) for _ in range(count)]
    start = time.time()
    sorted(nums, cmp=custom_cmp)
    cmp_duration = time.time() - start
    start = time.time()
    sorted(nums, key=custom_key)
    key_duration = time.time() - start
    print('%7d%20.2f%20.2f' % (count, cmp_duration, key_duration))

在我的笔记本上一次运行结果如下

  count        cmp_duration        key_duration
   1000                0.00                0.00
  10000                0.02                0.01
 100000                0.34                0.11
1000000                4.75                1.85

可以看到,当列表中数字的数量超过 100000 的时候,使用 key 函数的性能优势就非常明显了,比 cmp 快了 2~3 倍。

对于熟悉 Java 或 C++ 等其他编程语言的同学来说,可能更熟悉 cmp 的比较方式。其实 Python 3 中也可以通过 functools 工具包中的 cmp_to_key() 函数来将 cmp 转换成 key,从而使用接收两个参数的自定义比较函数 cmp。

import functools

num_list = [4, 2, 8, -9, 1, -3]

def custom_cmp(x, y):
    return x ** 2 - y ** 2

sorted_num_list = sorted(num_list, key=functools.cmp_to_key(custom_cmp))
print(sorted_num_list)

那么,cmp_to_key() 函数是如何将 cmp 转换成 key 的呢,我们可以通过源码一探究竟

def cmp_to_key(mycmp):
    """Convert a cmp= function into a key= function"""
    class K(object):
        __slots__ = ['obj']
        def __init__(self, obj):
            self.obj = obj
        def __lt__(self, other):
            return mycmp(self.obj, other.obj) < 0
        def __gt__(self, other):
            return mycmp(self.obj, other.obj) > 0
        def __eq__(self, other):
            return mycmp(self.obj, other.obj) == 0
        def __le__(self, other):
            return mycmp(self.obj, other.obj) <= 0
        def __ge__(self, other):
            return mycmp(self.obj, other.obj) >= 0
        __hash__ = None
    return K

其实 cmp_to_key() 返回的是一个类 K,只不过在类 K 中重载了各种比较运算符,重载的过程中使用到了自定义的比较函数 mycmp,使得 K 的大小比较逻辑与 mycmp 一致。这样,对于 num_list 中的每个元素 num 都会执行一次 K(num) 生成一个类 K 的实例,然后通过比较不同 K 的实例的大小进行排序。

虽然通过 cmp_to_key() 可以调用自定义的 cmp 函数,但是还是要优先使用 key 函数,因为通过 cmp_to_key() 方式会在排序过程中创建很多类 K 的实例,对性能有很大影响,下面是 cmp_to_key() 和 key 的性能比较

  count          cmp_to_key        key_duration
   1000                0.01                0.00
  10000                0.10                0.01
 100000                1.36                0.09
1000000               16.89                1.13

当 num_list 中的数量为 1000000 的时候 key 比 cmp_to_key 快了将近 15 倍。

本文主要介绍了如何在 sorted 函数中自定义比较逻辑,Python 2 中可以通过 cmp 或 key 来实现,cmp 接收 2 个参数,通过返回的数值来判断两个参数的大小,key 重新计算一个新的结果参与比较。在 Python 3 中,考虑到 cmp 的性能和冗余的原因,将其移除了。在 Python 3.2 中提供了 functools.cmp_to_key 这个函数来使用自定义的比较函数 cmp,但是出于性能的考虑,我们还是要优先使用 key 来进行排序。