""" classical priority queue

 it supports an arbitrary comparison function
 and is good for arbitrary length queues.

Based on pq2.py, optimized (about 3-4x faster).
"""

from bisect import insort

class PQ0:
    """priority queue using insertion sorting (bisect)

       This seems to be the fastest way, unless you want to use
       a different comparison metric, or unless the set grows *very* large.
    """
    def __init__(self, comparison = cmp):
        if comparison!=cmp:
            raise ValueError, "only cmp comparison supported for PQ0"
        self.data = []

    def empty(self):
        return len(self.data)==0

    def addelt(self, priority, elt):
        item = (priority, elt)
        insort(self.data, item)

    def largestp(self):
        return self.data[-1]

    def poplargest(self):
        data = self.data
        item = data[-1]
        del data[-1]
        return item

    def popsmallest(self):
        data = self.data
        item = data[0]
        del data[0]
        return item

class PQueue:
    """basic priority queue, using classical heap structure"""

    def __init__(self, comparison = cmp):
        self.cmp = comparison
        self.data = [None] * 8 # presize for 8 elements
        # first element is always empty, first used is 1, no data yet
        self.free = 1

    def empty(self):
        """empty test"""
        return self.free == 1

    def addelt(self, priority, elt):
        """add element by decreasing priority"""
        index = self.free
        try:
            self.data[ index ] = (priority, elt)
        except IndexError:
            # self.data is too small, double its size
            length = len( self.data )
            newdata = [ None ] * (2 * length)
            # store the old values
            newdata[:length] = self.data
            self.data = newdata
            # now there ought to be room!
            self.data[ index ] = (priority, elt)
        self.free = index + 1
        return self._checkposition(index)

    def popsmallest(self):
        """get/pop element with smallest priority"""
        if self.free < 2:
           raise IndexError, "priority queue is empty"
        smallest = self.data[1]
        self._removeentry(1)
        return smallest

    def _removeentry(self, index):
        """internal: remove an entry"""
        # restructure the "tree" if other elements remain
        free = self.free
        data = self.data
        last = free - 1
        if last > index:
           data[index] = data[last]
           data[last] = None
           self.free = last
           self._checkposition(index)
        else:
           data[index] = None
           self.free = last

    def smallestp(self):
        """get (priority, element) for smallest priority"""
        return self.data[1]
        
    # make sure a position has contents larger than parents,
    # and smaller than children, and if not, do a swap
    # and check the new position recursively...
    #
    def _checkposition(self, index):
        data = self.data
        comparison = self.cmp
        thisitem = data[index]
        thispriority = thisitem[0]
        free = self.free
        if index>=2:
            # check parent, possibly bubble up
            parent = index/2
            parentitem = data[parent]
            parentpriority = parentitem[0]
            if comparison(parentpriority, thispriority)>0:
                while 1:
                    data[parent] = thisitem
                    data[index] = parentitem
                    index = parent
                    if index<2:
                       return index
                    parent = index/2
                    parentitem = data[parent]
                    parentpriority = parentitem[0]
                    if comparison(parentpriority, thispriority)<=0:
                       return index
        # otherwise, check children
        # find the highest priority child:
        while 1:
            thechild = index*2
            if thechild>=free: return index
            childitem = data[thechild]
            childpriority = childitem[0]
            otherchild = thechild+1
            if otherchild<free:
                otherchilditem = data[otherchild]
                otherchildpriority = otherchilditem[0]
                if comparison(otherchildpriority, childpriority)<0:
                    thechild = otherchild
                    childitem = otherchilditem
                    childpriority = otherchildpriority
            # thisitem should be larger than childitem
            if comparison(thispriority, childpriority)<=0:
                return index
            data[index] = childitem
            data[thechild] = thisitem
            index = thechild

    def displaytree(self, index=1, indentlevel=0):
        print "   "*indentlevel, self.data[index], "at", index
        free = self.free
        for child in (index*2, index*2+1):
            if child<free:
               self.displaytree(child, indentlevel+1)

# this mixin must be combined with a "base" priority queue
# implementation.  It groups elements with common priorities.
# It should precede the "base" implementation in the inheritance
# hierarchy.
#
class PQEquivMixin:

   # requires a Q_implementation member

   # this initialization function MUST be called on subclass init.
   def __init__(self, comparison = cmp):
       # initialize self as a self.Q_implementation
       self.Q_implementation.__init__(self, cmp)
       # add a dictionary for holding elements at equivalent priorities
       self.EquivDict = {}

   def addelt(self, priority, elt):
       # is there a class of elements at this priority?
       EquivDict = self.EquivDict
       try:
           list = EquivDict[priority]
           list.append(elt)
       except KeyError:
           # there is none: add this element as a new priority group
           # First: add the priority to the queue
           self.Q_implementation.addelt(self, priority, priority)
           # Then: add the new group to the dictionary
           EquivDict[priority] = [elt]

   def smallest_group_p(self):
       # find the smallest priority
       (priority, dummy) = self.Q_implementation.smallestp(self)
       # return the priority and the associated group
       group = self.EquivDict[priority]
       return (priority, group)

   def smallest_p(self):
       (priority, group) = self.smallest_group_p()
       return (priority, group[-1])

   def popsmallest(self):
       (priority, group) = self.smallest_group_p()
       chosen = group[-1]
       del group[-1]
       # if group is now empty then delete this priority
       if group == []:
          del self.EquivDict[priority]
          self.Q_implementation.popsmallest(self)
       return (priority, chosen)

   def popsmallest_group(self):
       (priority, group) = self.smallest_group_p()
       self.Q_implementation.popsmallest(self)
       return group

# using the mixin:
class PQEquivBig(PQEquivMixin, PQueue):
   Q_implementation = PQueue