heap sort demo

heap_sort_demo.py

import unittest


def sift_down(lst, top, end):
    pos = end
    newitem = lst[pos]
    while pos > top:
        parentpos = (pos - 1) >> 1
        parent = lst[parentpos]
        if newitem < parent:
            lst[pos] = parent
            pos = parentpos
            continue
        else:
            break
    lst[pos] = newitem


def sift_up(lst, top):
    pos, end = top, len(lst)
    newitem = lst[pos]
    childpos = (pos << 1) + 1
    while childpos < end:
        rightpos = childpos + 1
        if rightpos < end and lst[rightpos] < lst[childpos]:
            childpos = rightpos
        lst[pos] = lst[childpos]
        pos = childpos
        childpos = (pos << 1) + 1
    lst[pos] = newitem
    sift_down(lst, top, pos)


def heapify(lst):
    for i in reversed(range(len(lst) >> 1)):
        sift_up(lst, i)


def heap_pop(lst):
    lastitem = lst.pop()
    if lst:
        returnitem = lst[0]
        lst[0] = lastitem
        sift_up(lst, 0)
        return returnitem
    return lastitem


def heap_push(lst, item):
    lst.append(item)
    sift_down(lst, 0, len(lst)-1)


def heap_sort(lst):
    heapify(lst)
    res = []
    while lst:
        res.append(heap_pop(lst))
    return res


class HeapqTest(unittest.TestCase):
    def setUp(self):
        import random
        self.heap = random.sample(range(0, 100), 15)

    def test_heapify(self):
        print('before sort: ', self.heap)
        sorted_heap = sorted(self.heap)
        heap_sorted_heap = heap_sort(self.heap)
        print('after sort:', sorted_heap)
        print('heap sort:', heap_sorted_heap)
        self.assertListEqual(sorted_heap, heap_sorted_heap)

    def test_heap_push(self):
        print('before sort:', self.heap)
        sorted_heap = sorted(self.heap)
        print('after sort:', sorted_heap)
        heap = []
        for i in self.heap:
            heap_push(heap, i)
        res = []
        while heap:
            res.append(heap.pop())
        print('heap sort:', res)

    def tearDown(self):
        self.heap = None


if __name__ == '__main__':
    unittest.main()
```

Result:

```
before sort: [66, 16, 2, 86, 55, 12, 64, 32, 8, 71, 22, 77, 10, 25, 74]
after sort: [2, 8, 10, 12, 16, 22, 25, 32, 55, 64, 66, 71, 74, 77, 86]
heap sort: [74, 64, 16, 77, 66, 71, 55, 86, 25, 12, 22, 32, 10, 8, 2]
.before sort:  [84, 87, 42, 30, 79, 55, 2, 71, 1, 91, 66, 70, 5, 75, 54]
after sort: [1, 2, 5, 30, 42, 54, 55, 66, 70, 71, 75, 79, 84, 87, 91]
heap sort: [1, 2, 5, 30, 42, 54, 55, 66, 70, 71, 75, 79, 84, 87, 91]
.
----------------------------------------------------------------------
Ran 2 tests in 0.007s

OK
```

在最小堆上获取nlargest元素的操作。

def nlargest(n, iterable):
    idx = 0
    heap = [0 for i in range(n)]
    for item in iterable:
        if idx < n:
            heap[idx] = item
        elif idx == n:
            heapify(heap)
            heap[0] = item
            sift_up(heap, 0)
        elif idx > n:
            if item < heap[0]:
                continue
            else:
                heap[0] = item
                sift_up(heap, 0)
        idx += 1
    res = [0 for i in range(n)]
    for i in reversed(range(n)):
        res[i] = heap_pop(heap)
    return res

class HeapqTest(unittest.TestCase):
    def setUp(self):
        import random
        self.heap = random.sample(range(0, 100), 15)

    def test_nlargest(self):
        print('before sort:', self.heap)
        sorted_heap = sorted(self.heap, reverse=True)
        print('after sort:', sorted_heap)
        largest7 = sorted_heap[:7]
        largest7_heap = nlargest(7, self.heap)
        print('7 largest via sort:', largest7)
        print('7 largest via heap sort:', largest7_heap)
        self.assertListEqual(largest7, largest7_heap)

    def tearDown(self):
        self.heap = None

before sort: [39, 28, 34, 90, 44, 45, 92, 30, 13, 74, 17, 40, 91, 23, 43]
after sort: [92, 91, 90, 74, 45, 44, 43, 40, 39, 34, 30, 28, 23, 17, 13]
7 largest via sort: [92, 91, 90, 74, 45, 44, 43]
7 largest via heap sort: [92, 91, 90, 74, 45, 44, 43]