如何使heapq根据特定属性来评估堆?

108

我希望可以使用堆(heap)来处理一组对象,而不仅仅是数字。这些对象会带有一个整数属性(attribute),通过该属性,堆可以进行排序。在Python中,使用heapq是最简单的方法,但是如果我想要按照特定属性(attribute)对其进行排序,应该怎么做呢?


这里大部分的回答已经过时了(2010年)。我添加了一个新的答案(2022年),适用于Python 3,并使用现代工具。它还解决了Python 3元组比较可能失败的问题,如果一些字段不支持比较(例如函数)。 - Raymond Hettinger
9个回答

151

根据文档中的示例,您可以使用元组,并且它将按照元组的第一个元素进行排序:

>>> h = []
>>> heappush(h, (5, 'write code'))
>>> heappush(h, (7, 'release product'))
>>> heappush(h, (1, 'write spec'))
>>> heappush(h, (3, 'create tests'))
>>> heappop(h)
(1, 'write spec')

所以如果你不想(或者不能)编写一个__cmp__方法,你可以在推送时手动提取排序关键字。

请注意,如果一对元组的第一个元素相等,那么将比较进一步的元素。如果这不是你想要的,请确保每个第一个元素都是唯一的。


43
请注意,如果一对元组的第一个元素相等,则会比较其他元素。由于文档不够清晰,建议将此内容加粗。您曾经假设在优先级相同的情况下,它会返回找到的第一个对象(这种假设没有充分理由,因此是您的责任)。 - JD Gamboa
好的观点。如果您插入一个元组,即(数字,字典),它不知道如何评估字典。 - Fred Guth
7
如果你有一个类似 (some_value, dict) 的元组,你可以将 (some_value, counter, dict) 插入堆中,以便在两个元组的 some_value 相等时使用递增计数器来打破平局。 - akshaynagpal
1
这个例子对我没用。有什么建议吗? lst = [(18, [3, 3]), (26, [5, -1]), (20, [-2, 4])] heapq.heapify(lst) - sanjay patel

110

heapqlist.sort 排序对象的方式相同,因此只需在类定义中定义一个方法 __cmp__(),该方法将自身与同一类的另一个实例进行比较:

def __cmp__(self, other):
    return cmp(self.intAttribute, other.intAttribute)

适用于 Python 2.x。

在 3.x 中请使用:

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

16
在Python 3.x中,__cmp__已经被移除,请使用__lt__代替。 - Daniel Stutzbach
13
在Python 2中,__lt__同样适用,因此最好完全避免使用__cmp__ - Daniel Stutzbach
20
就像你可以告诉任何一种排序方法基于对象的自然排序以外的标准来进行排序(例如 sortcmpkey),你应该能够告诉 heapq 基于不同的键来进行排序。换句话说,你不应该重新定义对象本身就能改变包含它的特定数据结构;你应该只需告诉数据结构本身即可。这是 heapq API 中一个值得注意的基本缺失部分。 - Glenn Maynard
3
大家为什么都要求使用__lt__而不是__gt__?这真的很重要吗? - Shrikant Shete
如果有时我想按此属性排序,有时又想按另一个属性排序怎么办? - jallen0927

42
根据官方文档,解决这个问题的方法是将条目存储为元组(请参阅第8.4.1节和8.4.2节)。 例如,在元组格式中,您的对象如下所示:(键,值1,值2) 当将对象(即元组)放入中时,它将使用对象中的第一个属性(在本例中为)进行比较。如果出现平局,则堆将使用下一个属性(即value_1)等等。 例如:
import heapq

heap = []
heapq.heappush(heap, (0,'one', 1))
heapq.heappush(heap, (1,'two', 11))
heapq.heappush(heap, (1, 'two', 2))
heapq.heappush(heap, (1, 'one', 3))
heapq.heappush(heap, (1,'two', 3))
heapq.heappush(heap, (1,'one', 4))
heapq.heappush(heap, (1,'two', 5))
heapq.heappush(heap, (1,'one', 1))

show_tree(heap)

输出:

                                      (0, 'one', 1)                                       
                (1, 'one', 1)                                (1, 'one', 4)                
    (1, 'one', 3)         (1, 'two', 3)         (1, 'two', 2)         (1, 'two', 5)     
(1, 'two', 11)

关于在python中漂亮地打印堆(更新了链接):show_tree()


15

Python 3 升级

其他答案已过时:

  • 有的是针对 Python 2 的。现在已经不存在 __cmp__ 方法。
  • 有些没有反映最佳实践,只针对 __lt__,而不是像 PEP 8 建议的所有 rich comparisons。
  • 有些没有使用现代工具,例如dataclassesattrgettertotal_ordering

使用 Dataclasses 的现代解决方案

使用dataclasses,可以轻松创建一个带有自定义排序的数据容器。例如,这是一个类,它在比较顺序中排除了姓名字段:

from dataclasses import dataclass, field

@dataclass(order=True)
class Person:
    name: str = field(compare=False)
    age: int

actors = [
    Person('T Hanks', 65),
    Person('E Olson', 33),
    Person('A Tapping', 58),
]

这在使用堆时完美运作:
>>> heapify(actors)
>>> heappop(actors)
Person(name='E Olson', age=33)
>>> heappop(actors)
Person(name='A Tapping', age=58)
>>> heappop(actors)
Person(name='T Hanks', age=65)

处理现有类

有时候你需要按照提供的数据来工作,并且需要控制比较的顺序,而不改变原始类。

解决方案是添加一个带有新比较方式的包装器。这样可以保留原始数据和它的类不变。以下是一个现代化的添加这样一个包装器的方法:

from functools import total_ordering
from operator import attrgetter

def new_compare(*field_names):
    extract = attrgetter(*field_names)
    @total_ordering
    class ComparisonWrapper:
        def __init__(self, obj):
            self.obj = obj
        def __eq__(self, other):
            return extract(self.obj) == extract(other.obj)
        def __lt__(self, other):
            return extract(self.obj) < extract(other.obj)
    return ComparisonWrapper

例如,您可能会收到以下数据,不能直接更改它或其类别:
class Person:
    def __init__(self, name, age):
        self.name = name
        self.age = age
    def __repr__(self):
        return f'Person({self.name!r}, {self.age})'

actors = [
    Person('T Hanks', 65),
    Person('E Olson', 33),
    Person('A Tapping', 58),
]

可以使用map()方法来优雅地应用这个包装器。要解除数据的包装,访问obj属性:

>>> from heapq import heapify, heappop

>>> data = list(map(new_compare('age'), actors))
>>> heapify(data)
>>> heappop(data).obj
Person('E Olson', 33)
>>> heappop(data).obj
Person('A Tapping', 58)
>>> heappop(data).obj
Person('T Hanks', 65)

包装器与元组装饰

正如现代文档中所指出的,传统的元组装饰方法在某些关键用例中已经不再适用。特别是,如果堆中的对象是函数,则形如(priority, task)的元组在Python 3中已经不再适用,因为函数无法进行比较。

新的建议是使用包装器,例如:

from dataclasses import dataclass, field
from typing import Any

@dataclass(order=True)
class PrioritizedItem:
    priority: int
    item: Any=field(compare=False)

即使item对象不可比较,也总是有效的。


11

我认为最简单的方法是覆盖heapq模块中现有的cmp_lt函数。以下是一个简短的示例:

import heapq

# your custom function. Here, comparing tuples a and b based on their 2nd element
def new_cmp_lt(self,a,b):
    return a[1]<b[1]

#override the existing "cmp_lt" module function with your function
heapq.cmp_lt=new_cmp_lt

#Now use everything like normally used
注意:如有更合适的编码实践建议,请有资格的人员进行评论。但是在编码面试中,如果时间有限且有更多任务需要完成而不是花时间正确地进行子类化,则仍然可能很有用。

注:更合格的人应该评论是否与推荐的编码实践冲突。但对于一些“快而脏”的东西(例如,在编写面试时时间有限且有更多事情要做而不是花时间正确地进行子类化),它仍然可以很有用。


9

我有同样的问题,但以上答案都没有完全解决我的问题,虽然有些接近但不够详细。无论如何,我进行了一些研究并尝试了这段代码,希望这对下一个正在寻找答案的人有所帮助:

使用元组的问题在于它只使用第一个项目,不是很灵活。我想要类似于C++中std::priority_queue的东西,就像这样: std::priority_queue<pair<int, int>, vector<pair<int, int>>, comparator> pq; 其中我可以设计自己的比较器,在实际应用中更为常见。

希望下面的代码片段有所帮助: https://repl.it/@gururajks/EvenAccurateCylinders

import heapq
class PQNode:

    def __init__(self, key, value):
        self.key = key
        self.value = value

    # compares the second value
    def __lt__(self, other):
        return self.value < other.value

    def __str__(self):
        return str("{} : {}".format(self.key, self.value))

input = [PQNode(1, 4), PQNode(7, 4), PQNode(6, 9), PQNode(2, 5)]
hinput = []
for item in input:
    heapq.heappush(hinput, item)

while (hinput):
    print (heapq.heappop(hinput))

我试了一下你的代码,在我的电脑上可以运行。我使用的是Python 3.6.5版本。我很好奇heappush()函数是如何进行比较的。这是通过PQNode类中特殊的_lt_()方法内在地完成的吗?如果没有它,这个程序肯定会崩溃,并显示编译器消息:Traceback (most recent call last): File "heap_example.py", line 18, in <module> heapq.heappush(hinput, item) TypeError: '<' not supported between instances of 'PQNode' and 'PQNode'。幸运的是,似乎_lt_()方法起到了作用,因为它现在可以正常工作。 - penguin2718
是的,PQNode 中的 __lt__ 是关键。请参见:https://docs.python.org/3/library/operator.html#operator.__lt__ - Arun

4

很遗憾,您不能这样做,尽管这是一个经常被请求的功能。

一种选择是将(键,值)元组插入堆中。但是,如果在比较值时抛出异常(在键相同时会进行比较),则此方法无法使用。

第二个选择是在类中定义__lt__(小于)方法,该方法将使用适当的属性来比较元素以进行排序。但是,如果对象是由另一个包创建的,或者您需要它们在程序的其他地方进行不同的比较,则可能无法实现。

第三个选择是使用sortedlist类来自blist模块(免责声明:我是作者)。sortedlist的构造函数接受一个key参数,该参数允许您指定一个函数来返回元素的排序键,类似于list.sortsortedkey参数。


我删除了之前的评论,因为我的blist问题可能是PEBCAK(再次感谢您的模块),所以我只复制了先前评论的第一部分:始终可以通过子类化或封装来定义具有__lt__的类。 - tzot

0

有一个名为heaps的模块。Github地址是https://github.com/gekco/heapy。您可以在类的实例化或从数组创建堆时应用自己的键/排序函数,这非常有用,因为这样可以节省每次执行操作时添加它作为参数的时间。

以下是一个示例,其中我希望列表中元组的最小元素位于堆的顶部:

>>> from heapy.heap import Heap 
>>> a = [(3, 5, 10), (-5, 3, 8), (7, 8, 9), (-4, 0, 2)]
>>> x = Heap.from_array(a, key=lambda t : t[-1])
>>> x.length
4
>>> x.top()
(-4, 0, 2)
>>> x.insert((-1, 0, 1))
>>> x.length
5
>>> x.top()
(-1, 0, 1)
>>> a
[(3, 5, 10), (-5, 3, 8), (7, 8, 9), (-4, 0, 2)]
 

0
你可以实现一个heapdict。注意使用popitem()来获取最低优先级的项。
import heapdict as hd
import string
import numpy as np

h = hd.heapdict()
keys = [char for char in string.ascii_lowercase[:10]]
vals = [i for i in np.random.randint(0,10, 10)]
for k,v in zip(keys,vals):
    h[k] = v
for i in range(len(vals)):
    print h.popitem()

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接