Is it possible to pass a comparator to a PriorityQueue in python

Question

I want to use a PriorityQueue to which I need to add objects of a class (say Node) which I cannot modify. I need these objects prioritized based on a field of the object.

I tried adding them as tuples(node.val, node) to the queue. It still gives me "TypeError: '<' not supported between instances of 'Node' and 'Node'". It still compares the 'node' object to resolve the ties.

a = Node(3)
b = Node(4)

q = queue.PriorityQueue()
q.put(a)
q.put(b) // Errors TypeError: '<' not supported between instances of 'Node' and 'Node'

q.put((a.val, a))
q.put((b.val, b)) // Errors TypeError: '<' not supported between instances of 'Node' and 'Node'

I understand that the cause of this error is that the Node objects need to be comparable to one another, and the way to achieve this according to the documentation is to at least implement lt and eq methods for Node. https://docs.python.org/3.1/library/stdtypes.html#comparisons.

But for my problem, since I wont be able to modify the Node class, will I be able to pass a way( a lambda or a Comparator) to the PriorityQueue to help it determine the ordering.(I am expecting something Java like)

q = PriorityQueue(comparator)

Any alternative to achieve this is also appreciated.(remember Node is not modifiable)

Answer 1

One solution, without passing a comparator, is to wrap your Node objects in another object, say ComparableNode, which implements the comparisons you would like on the Node object.

Let's assume you have this setup:

class Node:
    def __init__(self, val):
        self.val = val

    def __repr__(self):
        return "Node(val={})".format(self.val)


a = Node(3)
b = Node(4)
c = Node(4)

But you can't modify Node. So you can wrap it in a class that you create.

@functools.total_ordering
class ComparableNode(Node):
    def __gt__(self, other):
        return self.val > other.val

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


aa = ComparableNode(3)
bb = ComparableNode(4)
cc = ComparableNode(4)

or if you want to just pass Node objects that exist already:

@functools.total_ordering
class ComparableNode:
    def __init__(self, node):
        self.node = node

    def __gt__(self, other):
        return self.node.val > other.node.val

    def __eq__(self, other):
        return self.node.val == other.node.val

    def __repr__(self):
        return repr(self.node)


aa = ComparableNode(a)
bb = ComparableNode(b)
cc = ComparableNode(c)

Then just add them normally:

p = PriorityQueue()
p.put(aa)
p.put(bb)
p.put(cc)

p.queue  # [Node(val=3), Node(val=4), Node(val=4)]

Another way, if you can specify the priority of each element as you push it in. This is possible with heapq for example.

With heapq:

import heapq

q = []
heapq.heappush(q, (a.val, a))
heapq.heappush(q, (b.val, b))

q  # [(3, Node(val=3)), (4, Node(val=4))]

With queue.PriorityQueue, use tuples as suggested by the doc

from queue import PriorityQueue

p = PriorityQueue()
p.put((a.val, a))
p.put((b.val, b))

p.queue  # [(3, Node(val=3)), (4, Node(val=4))]

If you have duplicate values, then the comparison will look at the second element of the tuple, and will break because your Nodes are not comparable. In that case, it depends how you want to deal with duplicate values. If you want to preserve insert order in case of duplicates, you can do this:

from queue import PriorityQueue
from itertools import count

p = PriorityQueue()
index = count(0)

p.put((b.val, next(index), b))
p.put((c.val, next(index), c))

p.queue  # [(4, 0, Node(val=4)), (4, 1, Node(val=4))]

And proceed similarly with heapq. You could easily wrap this behavior in a small function so that you repeat yourself less.

p = PriorityQueue()
index = count(0)

def put_priority(node):
    p.put((node.val, next(index), node))

Answer 2

The PriorityQueue class does not allow for a key parameter, but you can subclass it to implicitly wrap each item inside a comparator object, here called a _Wrapper.

from queue import PriorityQueue


class _Wrapper:
    def __init__(self, item, key):
        self.item = item
        self.key = key

    def __lt__(self, other):
        return self.key(self.item) < other.key(other.item)

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


class KeyPriorityQueue(PriorityQueue):
    def __init__(self, key):
        self.key = key
        super().__init__()

    def _get(self):
        wrapper = super()._get()
        return wrapper.item

    def _put(self, item):
        super()._put(_Wrapper(item, self.key))

Usage

Here is an example of a priority queue which reverses the order or the elements.

key_pq = KeyPriorityQueue(key=lambda x: -x)

If you need a comparator and not a key, you can use the functools.cmp_to_key helper.

from functools import cmp_to_key

key_pq = KeyPriorityQueue(key=cmp_to_key(some_comparator))