Binary Heap

binary_heap.py

# Структура данных "Двоичная куча (Binary Heap)"
# https://habrahabr.ru/post/112222/


class Heap:
    def __init__(self, data=[]):
        # Функция построение кучи O(N)
        self.data = data
        for i in range(len(self.data) // 2, -1, -1):
            self.heapify(i)

    def __str__(self):
        # Служебная функция
        return str(' '.join(map(str, self.data)))

    def add(self, element):
        # Функция добавления нового элемента O(log N)

        # Добавляем новый элемент в конец кучи
        self.data.append(element)

        # Инициализируем индексы элемента и родителя
        i = len(self.data) - 1
        parent = (i - 1) // 2

        # Запускаем цикл по наведению порядка в куче
        # Буквально пока ребенок больше родителя, но при этом не выходя за границы массива,
        # ибо самый первый элемент всегда самый большой (встанет на -1 позицию = в конец)
        while i > 0 and self.data[i] > self.data[parent]:
            self.data[i], self.data[parent] = self.data[parent], self.data[i]
            i = parent
            parent = (i - 1) // 2

    def heapify(self, i):
        # Функция упорядочения двоичной кучи O(log N)

        # Бесконечный цикл (цикл с пост-условием), восстанавливающий порядок
        while True:
            # Инициализация основных переменных
            left_child = 2 * i + 1
            right_child = 2 * i + 2
            largest_child = i

            # Сравниваем верхушку с дочерними элементами, учитывая наличие всех детей
            if left_child < len(self.data) and self.data[left_child] > self.data[largest_child]:
                largest_child = left_child
            if right_child < len(self.data) and self.data[right_child] > self.data[largest_child]:
                largest_child = right_child

            # Если дети оказались меньше родителя – самое время тикать отсюда
            if largest_child == i:
                break

            # Выталкиваем больший из элементов наверх
            self.data[i], self.data[largest_child] = self.data[largest_child], self.data[i]
            # Переходим к "наибольшему"
            i = largest_child

    def pop_max(self):
        # Функция извлечения максимального элемента O(1)

        # Возвращаем нулевой элемент списка, а на его место ставим последний
        result, self.data[0] = self.data[0], self.data[len(self.data) - 1]
        del self.data[-1]
        return result


def heap_sort(a):
    # Функция сортировки с применением двоичной кучи O(N log N)
    c = Heap(a.copy())

    for i in range(len(a) - 1, -1, -1):
        a[i] = c.pop_max()
        c.heapify(0)