Python Samples - Priority Queue

From SwinBrain

The code described here is a minimal implementation of the set operations (and tests) required for algorithms like


import heapq
 
class PriorityQueue:
 
    # Helper class to manage the cost comparisons
    class Node:
        def __init__(self, data, cost):
            self.data = data
            self.cost = cost
 
        def __cmp__(self,rhs):
            return cmp(self.cost,rhs.cost)
        
    def __init__(self,defcost):
        self.heap = []
        self.defcost = defcost    
 
    def insert(self, data, cost):
        n = PriorityQueue.Node(data,cost)
        heapq.heappush(self.heap,n)
 
    def top(self):
        assert(not self.isEmpty())
        return self.heap[0].data
 
    def isEmpty(self):
        return self.heap == []
 
    def size(self):
        return len(self.heap)
 
    def topKey(self):
        cost = self.defcost
        if not self.isEmpty():
            cost = self.heap[0].cost
        return cost
 
    def pop(self):
        assert(not self.isEmpty())
        return heapq.heappop(self.heap).data
 
    def remove(self,data):
        i = -1
        for j in range(len(self.heap)):
            if self.heap[j].data == data:
                i = j
                break
        if i == -1:
            return
        self.heap = self.heap[:i]+self.heap[i+1:]
        heapq.heapify(self.heap)
 
 
if __name__ == "__main__":
    q = PriorityQueue(defcost = 1e3000)
    q.insert( 1, 1.0 )
    q.insert( 2, 0.9 )
    q.insert( 3, 1.1 )
    for n in range(4):
        print '---'
        print q.topKey()
        print q.top()
        print q.pop()