使用自定义比较谓词的heapq

Question

我正在尝试创建一个堆（heap），并且想用自定义的排序规则来排序。因为我放进去的值是“用户定义”的类型，所以我不能修改它们自带的比较规则。

有没有办法做到类似这样的事情：

h = heapq.heapify([...], key=my_lt_pred)
h = heapq.heappush(h, key=my_lt_pred)

更好的是，我可以把heapq的函数封装在我自己的容器里，这样就不需要一直传递排序规则了。

Answer 1

heapq的文档提到，堆里的元素可以是元组，元组的第一个元素是优先级，用来决定排序的顺序。

不过，更贴近你问题的是，文档里有一段讨论和示例代码，讲述了如何自己实现heapq的包装函数，以解决排序稳定性和优先级相同的元素处理等问题。

简单来说，他们的解决方案是让heapq中的每个元素都是一个三元组，包含优先级、一个入口计数和要插入的元素。这个入口计数可以确保优先级相同的元素按照它们被添加到heapq的顺序进行排序。

Answer 2

定义一个类，并重写 __lt__() 函数。下面是一个示例（在 Python 3.7 中有效）：

import heapq

class Node(object):
    def __init__(self, val: int):
        self.val = val

    def __repr__(self):
        return f'Node value: {self.val}'

    def __lt__(self, other):
        return self.val < other.val

heap = [Node(2), Node(0), Node(1), Node(4), Node(2)]
heapq.heapify(heap)
print(heap)  # output: [Node value: 0, Node value: 2, Node value: 1, Node value: 4, Node value: 2]

heapq.heappop(heap)
print(heap)  # output: [Node value: 1, Node value: 2, Node value: 2, Node value: 4]

Answer 3

根据heapq的文档，要自定义堆的顺序，我们需要把堆里的每个元素都放成一个元组，第一个元素是可以用普通的Python比较方式来比较的。

heapq模块里的函数有点麻烦，因为它们不是面向对象的，而且总是需要我们把堆对象（一个已经堆化的列表）作为第一个参数传进去。我们可以通过创建一个非常简单的包装类来一举两得，这样我们就可以指定一个key函数，并把堆当作一个对象来使用。

下面的这个类内部维护了一个列表，列表里的每个元素都是一个元组，第一个成员是一个关键值，这个关键值是在插入元素时根据key参数计算出来的，这个参数是在创建堆的时候传入的：

# -*- coding: utf-8 -*-
import heapq

class MyHeap(object):
    def __init__(self, initial=None, key=lambda x:x):
        self.key = key
        self.index = 0
        if initial:
            self._data = [(key(item), i, item) for i, item in enumerate(initial)]
            self.index = len(self._data)
            heapq.heapify(self._data)
        else:
            self._data = []

    def push(self, item):
        heapq.heappush(self._data, (self.key(item), self.index, item))
        self.index += 1

    def pop(self):
        return heapq.heappop(self._data)[2]

（额外的self.index部分是为了避免在比较关键值时出现平局的情况，而存储的值又无法直接比较——否则heapq可能会因为类型错误而失败）