Search code examples
pythonsortingpriority-queueheapq

heapq custom compareTo


I'm trying to define a custome method to make an ordered insert into a priority queue in python but not getting the expected results. Once defined the insert method into the queue like the following:

def insert(self, node):
    if isinstance(node, treeNode):
        heapq.heappush(self._frontier, (node._f, node))
    else:
        print("Error. It is not a node.")

And implementing the following lt in the 'node' class:

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

The insertion is not done by the 'f' attribute value which is what I wanna do, insert in an ascendant order determined by that value. Any help will be much appreciated.

Ex of failure:

[(141.09530289033592, <treeNode.treeNode object at 0x7f08bb840fd0>), (484.8315227978057, <treeNode.treeNode object at 0x7f08bb849da0>), (390.0514031446352, <treeNode.treeNode object at 0x7f08bb840e48>)]

It only puts in the first position de lowest value, which does make sense since using a priority queue but then the following ones are not sorted by the custom method I wanna declare.


Solution

  • As @Bakuriu mentioned the heapq only mantains the invariant of the heap, if you want to obtain the elements in order use nsmallest, for instance:

    import heapq
    
    
    class TreeNode:
        def __init__(self, f):
            self.f = f
    
        def __repr__(self):
            return 'f: {}'.format(self.f)
    
        def __lt__(self, other):
            return self.f < other.f
    
    
    class Frontier:
        def __init__(self):
            self.frontier = []
    
        def insert(self, node):
            if isinstance(node, TreeNode):
                heapq.heappush(self.frontier, (node.f, node))
            else:
                print("Error. It is not a node.")
    
        def __len__(self):
            return len(self.frontier)
    
    
    t = Frontier()
    
    for n in [TreeNode(141.09530289033592), TreeNode(484.8315227978057), TreeNode(390.0514031446352)]:
        t.insert(n)
    
    print(heapq.nsmallest(len(t), t.frontier))
    

    Output

    [(141.09530289033592, f: 141.09530289033592), (390.0514031446352, f: 390.0514031446352), (484.8315227978057, f: 484.8315227978057)]