Heapq with Custom Compare Predicate

heapq with custom compare predicate

According to the heapq documentation, the way to customize the heap order is to have each element on the heap to be a tuple, with the first tuple element being one that accepts normal Python comparisons.

The functions in the heapq module are a bit cumbersome (since they are not object-oriented), and always require our heap object (a heapified list) to be explicitly passed as the first parameter. We can kill two birds with one stone by creating a very simple wrapper class that will allow us to specify a key function, and present the heap as an object.

The class below keeps an internal list, where each element is a tuple, the first member of which is a key, calculated at element insertion time using the key parameter, passed at Heap instantiation:

# -*- 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]

(The extra self.index part is to avoid clashes when the evaluated key value is a draw and the stored value is not directly comparable - otherwise heapq could fail with TypeError)

How to make heapq evaluate the heap off of a specific attribute?

heapq sorts objects the same way list.sort does, so just define a method __cmp__() within your class definition, which will compare itself to another instance of the same class:

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

Works in Python 2.x.

In 3.x use:

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

How to heapify by field from custom objects

Add the magic __cmp__ comparison method to your class to avoid needing to do the tuple-decoration that Maksim describes:

>>> import heapq
>>> class MyObject(object):
...     def __init__(self, val):
...         self.val = val
...     def __cmp__(self, other):
...         return cmp(self.val, other.val)
...         
...     
... 
>>> q = []
>>> heapq.heappush(q, MyObject(50))
>>> heapq.heappush(q, MyObject(40))
>>> heapq.heappush(q, MyObject(30))
>>> heapq.heappush(q, MyObject(20))
>>> heapq.heappush(q, MyObject(200))
>>> obj = heapq.heappop(q)
>>> print obj.val
20

Note: Override __lt__ in Python 3, __cmp__ only in Python 2

Heapify list based on 2nd element of pair in python

You could append each element into the heap in whatever order you want. Heap recognizes them as (num1, num2) where num1 is used for ranking. Below is an option:

hp = []
for e in x:
    hp.append((e[1], e[0]))

heapq.heapify(hp)
new = []
while hp:
    y = heapq.heappop(hp)
    new.append([y[1], y[0]])
print(new)

The array new will have the order you want.

Heapq with equal priority

There's no reason you couldn't use the method you proposed: pop all items at the same instant, store them in a list, and then treat them sequentially. The basic idea is:

item = heap.pop()
itemlist.push(item)
while (heap not empty && heap.peek().priority == item.priority) {
    itemlist.push(heap.pop());
}

You'll want to convert that to real Python code, of course, but the basic idea works and is a perfectly valid use of a heap.

min heap in python

Yes, there is a way. Define a wrapping class that implements your custom comparator, and use a list of those instead of a list of your actual objects. That's about the best there is while still using the heapq module, since it provides no key= or cmp= arguments like the sorting functions/methods do.

def gen_wrapper(cmp):
    class Wrapper(object):
        def __init__(self, value): self.value = value
        def __cmp__(self, obj): return cmp(self.value, obj.value)
    return Wrapper

---

Related Topics