检查heapq是否包含值

Question

我正在使用heapq这个对象来存储我自己实现的类的对象。

import heapq

heap = []
element1 = Element('A', 1)
element2 = Element('B', 2)
element3 = Element('C', 3)
heapq.heappush(heap, element1)
heapq.heappush(heap, element2)
heapq.heappush(heap, element3)

在我的Element类中，我重写了__cmp__这个方法，以确保值是优先级。

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

现在我想写一个函数，检查堆中是否包含某个元素。比如说，如果我想检查element = Element('A', 1)是否在堆里，答案应该是True；如果我检查element = Element('A', 100)，答案也应该是True；但如果我想检查element = Element('D', 1)，答案就应该是False。我该如何实现这样的方法呢？有没有办法在不调用pop()方法的情况下检查heapq中的元素？

Answer 1

@enrico 提出的解决方案是可行的，它通过实现 __eq__ 方法来检查元素是否在堆中，同时用 __cmp__ 方法来确定元素的优先级。不过，这样做会产生一些奇怪的副作用。例如，Element('A', 1) 在某些情况下会同时被认为等于和小于 Element('A', 2)。

另外，你也可以直接使用普通的 tuples，而不是那个 Element 包装类。把数字放在第一位，元组的自然排序就能满足堆的需求。如果你想检查某个元素是否在堆中，可以用 zip 把这些项目组合起来，得到实际的键（在你的例子中就是字母），或者使用 any。

heap = []
heapq.heappush(heap, (1, 'A'))
heapq.heappush(heap, (3, 'C'))
heapq.heappush(heap, (2, 'B'))

print('A' in zip(*heap)[1])
print(any('D' == b for a, b in heap)

输出：

True
False

不过要注意的是（就像使用 __eq__ 和 in 一样），在最坏的情况下（也就是元素不在堆中），这会检查堆中的每一个元素。如果你在优先队列的场景中使用这种方法，会影响堆的 O(log n) 性能。因此，你可能需要使用另一个 set 或 dict 来存储和查找已经处理过的元素（及其对应的结果）和/或当前在堆中的元素。

Answer 2

在Element类里添加一个__eq__方法，这样你就可以用关键字in来检查元素是否存在了（如果没有__eq__，那么代码Element('A', 1) == Element('A', 1)会返回False）：

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

    def __eq__(self, other):
        return self.key == other.key

在Python中，堆其实就是列表，所以只需要使用下面的代码，__eq__会处理其他的事情：

Element('A', 1) in heap

示例

import heapq

heap = []
element1 = Element('A', 1)
element2 = Element('B', 2)
element3 = Element('C', 3)
heapq.heappush(heap, element1)
heapq.heappush(heap, element2)
heapq.heappush(heap, element3)

print Element('A', 1) in heap      # True
print Element('A', 100) in heap    # True
print Element('D', 1) in heap      # False