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 heapqclass 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)

Define a class, in which override the __lt__() function. See example below (works in Python 3.7):

import heapqclass Node(object):    def __init__(self, val: int):        self.val = val    def __repr__(self):        return f'Node value: {self.val}'    def __lt__(self, other):        return self.val < other.valheap = [Node(2), Node(0), Node(1), Node(4), Node(2)]heapq.heapify(heap)print(heap)  # output: [Node value: 0, Node value: 2, Node value: 1, Node value: 4, Node value: 2]heapq.heappop(heap)print(heap)  # output: [Node value: 1, Node value: 2, Node value: 2, Node value: 4]

The heapq documentation suggests that heap elements could be tuples in which the first element is the priority and defines the sort order.

More pertinent to your question, however, is that the documentation includes a discussion with sample code of how one could implement their own heapq wrapper functions to deal with the problems of sort stability and elements with equal priority (among other issues).

In a nutshell, their solution is to have each element in the heapq be a triple with the priority, an entry count and the element to be inserted. The entry count ensures that elements with the same priority a sorted in the order they were added to the heapq.