Python Heap

gistfile1.py

max_predicate = lambda a, b: a > b


def heap_swap(heap, a, b):
    """
    Swaps two elements in array
    """
    t = heap[a]
    heap[a] = heap[b]
    heap[b] = t


def heapify(heap, heap_size, x, predicate=max_predicate):
    """
    Makes a heap of subtree
    """
    left = 2 * (x + 1) - 1
    right = 2 * (x + 1)
    largest = x

    if left < heap_size and predicate(heap[left], heap[x]):
        largest = left
    if right < heap_size and predicate(heap[right], heap[largest]):
        largest = right
    if largest != x:
        heap_swap(heap, x, largest)
        heapify(heap, heap_size, largest, predicate)


def make_heap(heap, predicate=max_predicate):
    """
    Turns array into heap
    """
    heap_size = len(heap)
    for i in reversed(range(0, heap_size // 2)):
        heapify(heap, heap_size, i, predicate)


def heap_sort(array, predicate=max_predicate):
    """
    Sorts array
    """
    make_heap(array, predicate)
    for i in reversed(range(1, len(array))):
        heap_swap(array, 0, i)
        heapify(array, i, 0, predicate)


def push_heap(heap, value, predicate=max_predicate):
    """
    Adds new element into heap
    """
    heap.append(value)
    current = len(heap) - 1
    while current > 0:
        parent = (current - 1) // 2
        if predicate(heap[current], heap[parent]):
            heap_swap(heap, parent, current)
            current = parent
        else:
            break


def pop_heap(heap, predicate=max_predicate):
    """
    Deletes top element from heap
    """
    pop_value = heap[0]
    heap_swap(heap, 0, len(heap) - 1)
    heapify(heap, len(heap) - 1, 0, predicate)
    return pop_value


# heap = [4, 6, 1, 3, 10, 5, 11, 8, 9]
#
# heap_sort(heap)
# print(heap)
#
# heap_sort(heap, lambda a, b: a < b)
# print(heap)

heap = [1, 3, 2, 10, 5, 2]
make_heap(heap)
print(heap)

heap = []

minPred = lambda a, b: a < b

push_heap(heap, 1)
push_heap(heap, 3)
push_heap(heap, 2)
push_heap(heap, 10)
push_heap(heap, 5)
push_heap(heap, 2)
print(heap)

while len(heap):
    print(pop_heap(heap), end=' ')
    heap = heap[:-1]
print()