Check if heapq contains value

Question

I am using the heapq object to store objects of class which I implemented.

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)

In my class Element I overwrite method __cmp__ to ensure that value is the priority

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

Now I want to write a function, which checks if heap contains element, such that if I want to check if element = Element('A', 1) is in the heap, the answer will be True, if I will check element = Element('A',100) the answer will be also True, but if I want to check element = Element('D',1) the answer will be False. How can I implement such method? Is it possible to check elements of heapq without calling pop() method?

Answer 1

Add the method __eq__ in Element so you can check for membership using the keyword in (without __eq__ the code Element('A', 1) == Element('A', 1) would give False):

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

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

Heaps are just lists in python, so just use the following and __eq__ will do the rest:

Element('A', 1) in heap

Example

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