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