堆 python 实现

# -*- coding: UTF-8 -*-
"""
最大堆
"""

class MaxHeap:
    def __init__(self):
        self.array = []
        return

    def _parent(self, k):
        if k == 0:
            raise Exception('%s has no parent' % k)
        return (k-1)/2

    def _left_child(self, k):
        return k*2 + 1

    def _right_child(self, k):
        return k*2 + 2

    def _compare(self, i, j):
        if self.array[i] < self.array[j]:
            return -1
        elif self.array[i] == self.array[j]:
            return 0
        else:
            return 1

    def _swap(self, k, j):
        self.array[k], self.array[j] = self.array[j], self.array[k]

    def _swiftup(self, k):
        while k > 0 and self._compare(self._parent(k), k) < 0:
            self._swap(self._parent(k), k)
            k = self._parent(k)
        return

    def _swiftdown(self, k):
        while self._left_child(k) < len(self.array):
            lc = self._left_child(k)
            rc = self._right_child(k)
            j = lc
            if rc < len(self.array) and self._compare(lc, rc) < 0:
                j = rc
            if self._compare(j, rc) >= 0:
                break
            self._swap(j, k)
        return

    def add(self, value):
        self.array.append(value)
        self._swiftup(len(self.array)-1)
        return

    def remove_head(self):
        # 交换首尾
        self._swap(0, len(self.array)-1)
        del(self.array[len(self.array)-1])
        self._swiftdown(0)
        return

    def get_head(self):
        if len(self.array) == 0:
            print ("heap is empty!!!!")
            return 0
        return self.array[0]

    def get_heap_array(self):
        return self.array



def test_heap():
    arr = []
    for i in range(0, 100):
        r = random.randint(0, 200)
        arr.append(r)
    heap = MaxHeap()
    for i in range(len(arr)):
        heap.add(arr[i])

    print ("heap = %s" % heap.get_heap_array())
    print ("head = %s" % heap.get_head())
    return


if __name__ == '__main__':
    print "start"
    test_heap()