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?

Is `Element` just a wrapper for values and their priorities used only for this heap, or do you have to use this class for other reasons? — tobias_k, Aug 14 '14 at 20:16
can you add a bit more info as to what this class is actually for and is this for python 2 or 3? — Padraic Cunningham, Aug 14 '14 at 21:33
@Ziva In this case, you could take particular care that implementing `eq` and `cmp` in an inconsistent way does not cause problems in other places. — tobias_k, Aug 15 '14 at 05:48

enrico.bacis · Accepted Answer · 2014-08-14T20:23:20.270

4

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

edited Aug 14 '14 at 20:23

answered Aug 14 '14 at 20:15

enrico.bacis

30,497
10
86
115

1

I think `__eq__` should compare only the `key`, so that `(A, 1)` == `(A, 100)` – tobias_k Aug 14 '14 at 20:19
@tobias_k It's possible, but maybe he wants to allow duplicates, it depends on the semantic. Let's see what the OP says and I'll adapt the code if needed. Thanks for the observation! – enrico.bacis Aug 14 '14 at 20:21
I think it's pretty clear: "if I will check element = Element('A',100) the answer will be also True" – tobias_k Aug 14 '14 at 20:22
Doesn't `Element('A', 1) in heap` operate in `O(n)` time? A sorted binary tree should be searchable in `O(log n)` time. I was hoping to find a way to do this with the benefit of the `heapq` sorting. Unless I'm missing something? I think this SO post answers this question: https://stackoverflow.com/questions/49619369/find-the-position-of-an-element-in-a-heap – David Parks Oct 12 '22 at 23:11

tobias_k · Answer 2 · 2020-07-21T09:59:36.557

The solution by @enrico works, implementing __eq__ to check whether elements are in the heap, and __cmp__ for prioritizing the elements. However, it will produce some strange side effects. For example, Element('A', 1) will at the same time be == to and < than Element('A', 2).

Alternatively, you could just use regular tuples instead of that Element wrapper class. With the number in the first place, the tuples natural ordering will suffice for the heap, and for checking whether some element is in the heap, you could zip the items to get a list of the actual keys (in your example: the letters) or use 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)

Output:

True
False

Note, however, that (just like using __eq__ and in) this will in the worst case (the element is not in the heap) test each element in the heap, which, if this is used in sort of a priority-queue setting, will ruin the otherwise O(log n) performance of the heap. Instead, you might want to use another set or dict besides the heap to store and look up elements that have already been handled (and their corresponding result) and/or are currently in the heap.

Check if heapq contains value

2 Answers2

Example