Custom comparator for building PriorityQueue in Python

Question

Iam trying to build a priority queue using PriorityQueue in Python, but instead of element to be considered for priority comparison, I want it to use the return value from a function after passing the element to the function , similar to sorted(mtlist,key = myfun), is there a way to achieve this,

Answer 1

Rather than inserting your elements directly into the queue, wrap each element in a tuple, where the first element in the tuple is the desired sorting key. Tuples are sorted by in order of their elements (i.e., first element is compared first), hence why the sorting key needs to come first.

import heapq

queue = []
my_list = [...]
for element in my_list:
    heapq.heappush(queue, (my_func(element), element))

Answer 2

If you have a wrapper class for the elements, then you can use operator overloading.

For example, lets say you have a CustomNumber class (equivalent to your elements) where the order is determined by the modulo 16 value (the private function __f()), the you can override the comparison operators like:

class CustomNumber:
    def __init__(self, value):
        self.value = value

    def __f(self, x):
        return x % 16

    def __lt__(self, obj):
        """self < obj."""
        return self.__f(self.value) < self.__f(obj.value)

    def __le__(self, obj):
        """self <= obj."""
        return self.__f(self.value) <= self.__f(obj.value)

    def __eq__(self, obj):
        """self == obj."""
        return self.__f(self.value) == self.__f(obj.value)

    def __ne__(self, obj):
        """self != obj."""
        return self.__f(self.value) != self.__f(obj.value)

    def __gt__(self, obj):
        """self > obj."""
        return self.__f(self.value) > self.__f(obj.value)

    def __ge__(self, obj):
        """self >= obj."""
        return self.__f(self.value) >= self.__f(obj.value)

Such that the following code:

a = CustomNumber(16)
b = CustomNumber(14)

print('a < b =', a < b)
print('a <= b =', a <= b)
print('a == b =', a == b)
print('a != b =', a != b)
print('a > b =', a > b)
print('a >= b =', a >= b)

prints:

a < b = True
a <= b = True
a == b = False
a != b = True
a > b = False
a >= b = False