编程Python

在heapq中实现自定义比较逻辑

字数统计: 1.3k字   |   阅读时长: 4分
,

摘要: 在 heapq 中自定义比较函数

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

在文章 heapq的用法 中，我们介绍了在 Python 中用 heapq 维护最小堆的方法。但是 heapq 中默认是最小堆，也就是堆头始终是最小元素。

如果想要最大堆，一种方法是对所有的数取相反数，这样依然可以用 heapq 来维护。参考文章 在heqpq中维护大顶堆：取相反数的方式。

与 heapq 类似的组件还有一个 queue.PriorityQueue，但 queue.PriorityQueue 默认也是最小值优先，想要最大堆的话依然可以通过相反数的方式实现。

与相反数的方案相比，我们更关心的是在堆中如何自定义比较逻辑，因为自定义比较逻辑除了可以解决最大堆的问题，还可以解决在堆中维护自定义对象等更复杂的问题。

heapq 不支持自定义键函数

我们知道 Python3 在排序中自定义比较的时候，是可以通过自定义比较函数或者键函数来解决的，参考文章：

这种方法与 C++ 中自定义比较函数的方法类似，参考以下文章：

因此很自然地想，是不是能够类似地定义键函数，然后在创建堆或优先级队列的实例时指定自定义键函数。类似于下面这样：

1
2
3
4
5
6
7
8
9
10
11
12
13
class Key:
    def __init__(self, ...):
        # 初始化键函数持有的额外信息

    def __call__(self, i):
        # 对象 i 的键函数的逻辑

# 创建键函数实例
mykey = Key(...)
# 堆中的数据
data = [...]
# 创建堆
heapq.heapify(data, key=mykey)

但是很遗憾，heapq 和 queue.PriorityQueue 均不支持自定义比较器。也就是说上面代码中 heapq.heapify(data, key=mykey) 是脑补出来的，heapq 中没有 key 参数。

妥协的思路

Python3 在排序时支持指定键函数，在文章 Python3自定义排序：直接定义键函数，或先定义比较函数再转换为键函数 中有详细阐述。类似地，对堆中的对象设计键函数也有两种方法：

方法一是先定义 HeapCmp 类，在其 __call__ 方法中实现比较函数的逻辑，这样 HeapCmp 实例就相当于比较函数，然后再通过 cmp_to_key 转换为键函数即可。这样转换成的键函数输入原始对象，输出的可以理解为还是该对象本身，只是增加了自定义的 __lt__，在这个 __lt__ 中实际调用的就是由 HeapCmp 给出的比较函数。

方法二是定义 HeapKey 类，在其 __call__ 方法中实现键函数的逻辑，返回一个带有 __lt__ 的对象作为比较键。这样 HeapKey 的实例就相当于一个键函数，它输入原始对象，输出用于比较的键。

但是 Python 标准库的 heqpq 和 PriorityQueue 中并不直接支持指定键函数，于是就有两种妥协的思路：

如果定义了 HeapCmp 再通过 cmp_to_key 得到一个类包装器 mykey，用于给原对象包装一个 __lt__（逻辑来源于 HeapCmp）。这样可以直接压入 mykey(myobj) 在堆中维护，弹出 item 时，其 item.obj 就是原对象。

如果定义了 HeapKey 进而直接得到键函数 mykey，可以构造元组 (mykey(myobj), myobj)，但这样的话堆中如果有两个元组的第一位元素相同，还是会启动对第二位的原始对象的比较，而这里原始对象可能并没有 __lt__，这就产生问题。所以在Python3自定义排序中可行的直接定义键函数的方式在 heapq 这里不适用，根本的解法还是给原对象定义 __lt__。

最终解法

综上，在 Python3 中如果想在 heapq 和 PriorityQueue 中自定义比较逻辑的方法如下：

(1) 优先给原对象直接定义 __lt__，相当于 C++ 中的重载对象中的小于号，参考文章 自定义排序/最值/堆的比较规则：自定义对象的小于方法。

(2) 如果原对象是数值型这种基础类型但比较有另外的逻辑，或者比较时需要未知的额外信息，不方便直接给原对象定义 __lt__。可以先定义 HeapCmp，在其中放入额外信息和比较逻辑，然后通过 cmp_to_key 转换为一个类包装器 mykey，相当于给原对象包装了一个 __lt__。之后操作这个被包装后的对象。

例子

Share