带自定义比较谓词的 heapq

我尝试用自定义排序谓词构建堆。由于输入的值属于“用户定义”类型，因此我不能修改它们内置的比较谓词。

有没有一种方法可以这样做:

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

或者更好的是，我可以将 heapq函数封装在我自己的容器中，这样就不需要一直传递谓词了。

111621 次浏览
小开

Heapq 文档建议堆元素可以是元组，其中第一个元素是优先级，并定义排序顺序。

但是，与您的问题更相关的是，文档中包含了一个 与示例代码讨论，说明如何实现自己的 heapq 包装函数来处理排序稳定性和具有同等优先级的元素的问题(以及其他问题)。

简而言之，他们的解决方案是让 heapq 中的每个元素都是一个具有优先级、入口计数和要插入的元素的三元组。条目计数确保具有相同优先级的元素按其添加到 heapq 的顺序排序。

小开
最佳答案

根据 Heapq 文档，定制堆顺序的方法是让堆上的每个元素都是一个 tuple，第一个 tuple 元素是一个接受普通 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 可能会因 TypeError 而失败)

小开

这两种答案的局限性在于，它们都不允许把关系当作关系来对待。在第一种情况下，通过比较项目来打破关系，在第二种情况下，通过比较输入顺序来打破关系。让领带成为领带会更快，如果有很多领带，那么就会有很大的不同。基于上述和文档，目前还不清楚这是否可以在 heapq 中实现。Heapq 不接受密钥，而同一模块中从它派生的函数却接受密钥，这看起来确实很奇怪。
附言: 如果你按照第一条评论中的链接(“可能重复...”) ，还有另一个定义 le 的建议，它看起来像是一个解决方案。

小开

定义一个类，在这个类中覆盖 __lt__()函数:

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]
小开 
setattr(ListNode, "__lt__", lambda self, other: self.val <= other.val)

使用它来比较 heapq 中对象的值

小开

在 python3中，可以从 functools模块使用 cmp_to_key。

假设您需要一个由三个元素组成的优先级队列，并使用 last 属性指定优先级。

from heapq import *
from functools import cmp_to_key
def mycmp(triplet_left, triplet_right):
key_l, key_r = triplet_left[2], triplet_right[2]
if key_l > key_r:
return -1  # larger first
elif key_l == key_r:
return 0  # equal
else:
return 1




WrapperCls = cmp_to_key(mycmp)
pq = []
myobj = tuple(1, 2, "anystring")
# to push an object myobj into pq
heappush(pq, WrapperCls(myobj))
# to get the heap top use the `obj` attribute
inner = pq[0].obj

性能测试:

环境

Python 3.10.2

密码

from functools import cmp_to_key
from timeit import default_timer as time
from random import randint
from heapq import *


class WrapperCls1:
__slots__ = 'obj'
def __init__(self, obj):
self.obj = obj
def __lt__(self, other):
kl, kr = self.obj[2], other.obj[2]
return True if kl > kr else False


def cmp_class2(obj1, obj2):
kl, kr = obj1[2], obj2[2]
return -1 if kl > kr else 0 if kl == kr else 1


WrapperCls2 = cmp_to_key(cmp_class2)


triplets = [[randint(-1000000, 1000000) for _ in range(3)] for _ in range(100000)]
# tuple_triplets = [tuple(randint(-1000000, 1000000) for _ in range(3)) for _ in range(100000)]


def test_cls1():
pq = []
for triplet in triplets:
heappush(pq, WrapperCls1(triplet))
        

def test_cls2():
pq = []
for triplet in triplets:
heappush(pq, WrapperCls2(triplet))


def test_cls3():
pq = []
for triplet in triplets:
heappush(pq, (-triplet[2], triplet))


start = time()
for _ in range(10):
test_cls1()
# test_cls2()
# test_cls3()
print("total running time (seconds): ", -start+(start:=time()))

结果

每个函数使用 list而不是 tuple:

  • WrapperCls1:16.2 ms
  • 带 __slots__的 WrapperCls1:9.8 ms
  • WrapperCls2:8.6 ms
  • 将优先级属性移动到第一个位置(不支持 习俗谓词) : 6.0 ms。

因此，此方法略快于使用具有重写的 __lt__()函数和 __slots__属性的自定义类。


提交答案