Heap, Heapsort Python

heap.py

"""
Heap (Min Heap), Heapsort
Programming Pearls: column 14
"""
import sys
import random
import time

sys.setrecursionlimit(1500)


def swap(h, i, j):
    h[i], h[j] = h[j], h[i]


def insert(h, n):
    h.append(n)
    siftup(h, len(h) - 1)


def pop(h):
    if len(h) > 1:
        t = h[1]
        swap(h, 1, len(h) - 1)
        h.pop()
        siftdown(h, 1)
        return t


def siftup(h, i):
    while i >= 2:
        j = i // 2
        if h[j] <= h[i]:
            return
        swap(h, i, j)
        i = j


def siftdown(h, i):
    while i*2 < len(h):
        l = i*2
        r = l+1
        j = l if r >= len(h) or h[l] <= h[r] else r
        if h[i] <= h[j]:
            break
        swap(h, i, j)
        i = j


def siftdownn(h, i, n):
    # slightly modified version of siftdown
    # siftdown element of index i to at max index n-1
    while i*2 < n:
        l = i*2
        r = l+1
        j = l if r >= n or h[l] <= h[r] else r
        if h[i] <= h[j]:
            break
        swap(h, i, j)
        i = j


def heapsort1(lst):
    # require extra O(N) memory

    # 1-based heap is easy to code
    x = [-1]

    for e in lst:
        insert(x, e)
    r = []
    while True:
        t = pop(x)
        if t is None:
            break
        r.append(t)
    lst[:] = r


def heapsort2(lst):
    # in-place sort, no extra memory

    # 1-based heap is easy to code
    lst.insert(0, -1)  # O(N)

    n = len(lst)
    # build heap incrementally from lst[1..1], lst[1..2] to lst[1..N]
    for i in range(2, n):
        siftup(lst, i)

    # swap min with last element in array
    # maintain heap from lst[1..N], lst[1..N-1] to lst[1..1]
    for i in range(n-1, 1, -1):
        swap(lst, i, 1)
        siftdownn(lst, 1, i)

    lst.pop(0)  # O(N)

    # since we are using a min-heap, the smallest value is at the back
    # flip the whole array
    for i in range(len(lst) // 2):
        swap(lst, i, len(lst)-i-1)


for sort in [heapsort1, heapsort2]:
    for size in [100, 1000, 10000]:
        lists = [[random.randint(1, 100) for i in range(size)],
                 [random.uniform(1, 100) for i in range(size)],
                 [1] * size,
                 list(range(size)),
                 list(reversed(range(size)))]
        for i, lst in enumerate(lists):
            t = time.time()
            sort(lst)
            print('%s\t%-5d\t%d\t%.6f' % (sort, size, i, time.time() - t))
            assert lst == sorted(lst)
    print('---')