heaps.py

# Programming Pearls: column 14

class Heaps(object):
    """
    >>> h = Heaps()
    >>> len(h)
    0
    >>> h.extend([12, 20, 15, 3, 13, 23])
    >>> len(h)
    6
    >>> print h
    [3, 12, 15, 20, 13, 23]
    >>> h.extractmin()
    3
    >>> print h
    [12, 13, 15, 20, 23]
    >>> h.extractmin()
    12
    >>> print h
    [13, 20, 15, 23]
    >>> h.extractmin()
    13
    >>> print h
    [15, 20, 23]
    >>> h.extend([3, 2, 3])
    >>> print h
    [2, 3, 3, 20, 15, 23]
    >>> h.extractmin()
    2
    """

    def __init__(self):
        self.__items = []

    def __len__(self):
        return len(self.__items)

    def __str__(self):
        return str(self.__items)

    # value(i) = x[i]
    def __getitem__(self, i):
        return self.__items[i]

    def __setitem__(self, i, v):
        self.__items[i] = v

    def extend(self, items):
        for i in items:
            self.append(i)

    def append(self, item):
        self.__items.append(item)
        self.shiftup()

    def extractmin(self):
        if not self.__items:
            return None

        m = self[0]
        self[0] = self.__items.pop()
        self.shiftdown()
        return m

    # O(log n)
    def shiftup(self):
        i = len(self) - 1
        while True:
            if i is 0:
                break
            parent = ((i + 1) / 2) - 1
            if self[parent] <= self[i]:
                break
            self[parent], self[i] = self[i], self[parent]
            i = parent

    # O(log n)
    def shiftdown(self):
        i, n = 0, len(self)
        while True:
            child = (2 * (i + 1)) - 1
            if child >= n:
                break
            if child + 1 < n:
                if self[child + 1] < self[child]:
                    child += 1

            if self[i] <= self[child]:
                break
            self[child], self[i] = self[i], self[child]
            i = child


if __name__ == '__main__':
    import doctest
    doctest.testmod()