heapz.py

import random
import sys
import big_o

import heapq

def heapify(x: list):
    for i, value in reversed(tuple(enumerate(x))):
        while True:
            left_child_index = 2 * i + 1
            right_child_index = 2 * i + 2
            try:
                left_child_value = x[left_child_index]
            except IndexError:
                break
            try:
                right_child_value = x[right_child_index]
            except IndexError:
                if value > left_child_value:
                    x[i] = left_child_value
                    x[left_child_index] = value
                    i = left_child_index
                else:
                    break
            else:
                if value > right_child_value <= left_child_value:
                    x[i] = right_child_value
                    x[right_child_index] = value
                    i = right_child_index
                elif value > left_child_value <= right_child_value:
                    x[i] = left_child_value
                    x[left_child_index] = value
                    i = left_child_index
                else:
                    break


def heappop(x: list):
    result = x[0]
    try:
        x[0] = x.pop()
    except IndexError:
        # One-item heap.
        return result
    i = 0
    while True:
        value = x[i]
        left_child_index = 2 * i + 1
        right_child_index = 2 * i + 2
        try:
            left_child_value = x[left_child_index]
        except IndexError:
            return result
        try:
            right_child_value = x[right_child_index]
        except IndexError:
            if value > left_child_value:
                x[i] = left_child_value
                x[left_child_index] = value
                i = left_child_index
                continue
        else:
            if value > right_child_value <= left_child_value:
                x[i] = right_child_value
                x[right_child_index] = value
                i = right_child_index
                continue

            elif value > left_child_value <= right_child_value:
                x[i] = left_child_value
                x[left_child_index] = value
                i = left_child_index
                continue

        # If we reached this point, the children of the current node are bigger
        # than it, meaning the tree is now valid.
        return result


def test():
    for _ in range(10**6):
        if _ % (10 ** 4) == 0:
            print('.', end='')
            sys.stdout.flush()

        array = [random.randint(0, 20) for _ in range(random.randint(1, 20))]
        heapify(array)

        # copy = array[:]
        # heapq.heapify(copy)
        # assert copy == array

        sorted_array = sorted(array)
        result = []
        while array:
            result.append(heappop(array))

            copy = array[:]
            heapify(copy)
            assert copy == array
        assert result == sorted_array


def data_generator_heapify(n):
    return [random.randint(0, 20) for _ in range(n)]


def data_generator_heappop(n):
    array = data_generator_heapify(n)
    heapify(array)
    return array



best, others = big_o.big_o(heapify, data_generator_heapify, max_n=10**7)
print(best)
print()
for other in others:
    print(other)
print()


best, others = big_o.big_o(heappop, data_generator_heappop, max_n=10**7)
print(best)
print()
for other in others:
    print(other)
print()