Обзор

Куча - это структура данных для работы с приоритетами. Она позволяет быстро получать минимальный или максимальный элемент, не сортируя весь набор данных после каждого изменения.

В минимальной куче на вершине всегда лежит самый маленький элемент. В максимальной куче - самый большой.

Идея структуры

Куча поддерживает две основные операции:

1. Добавить новый элемент за O(log n).

2. Извлечь элемент с наивысшим приоритетом за O(log n).

Посмотреть на вершину можно за O(1).

Это удобно, когда нам не нужен полный порядок всех элементов. Нам нужен только следующий лучший элемент: самый маленький, самый большой или один из k лучших.

Реализация на массиве

На практике кучу почти всегда хранят в массиве. Дерево кучи - полное бинарное дерево, и его можно разложить по уровням без указателей: корень лежит в индексе 0, а для узла с индексом i работают формулы:

parent(i) = Math.floor((i - 1) / 2)
left(i)   = 2 * i + 1
right(i)  = 2 * i + 2

Например, для минимальной кучи:

       1
      / \
     3   2
    / \
   7   4
 
data = [1, 3, 2, 7, 4]

Свойство кучи простое: значение узла не больше значений его детей. В минимальной куче это значит, что минимум всегда в корне, то есть в data[0].

heapifyUp используют при добавлении элемента. Новое значение кладут в конец массива, а потом поднимают его вверх, пока оно меньше родителя. Так восстанавливается свойство кучи после push.

heapifyDown используют при извлечении минимума. Мы забираем корень, последний элемент массива переносим на его место и опускаем вниз, каждый раз меняясь с меньшим из двух детей. Так восстанавливается свойство кучи после pop.

Вот как может выглядеть реализация кучи на массиве:

data = [];
 
function parent(i) {
  return Math.floor((i - 1) / 2);
}
 
function left(i) {
  return 2 * i + 1;
}
 
function right(i) {
  return 2 * i + 2;
}
 
function swap(i, j) {
  temp = data[i];
  data[i] = data[j];
  data[j] = temp;
}
 
function push(value) {
  data.push(value);
  heapifyUp(data.length - 1);
}
 
function pop() {
  if (data.length == 0) return null;
  root = data[0];
  last = data.pop();
  if (data.length > 0) {
    data[0] = last;
    heapifyDown(0);
  }
  return root;
}
 
function heapifyUp(index) {
  while (index > 0) {
    p = parent(index);
    if (data[p] <= data[index]) break;
    swap(p, index);
    index = p;
  }
}
 
function heapifyDown(index) {
  while (true) {
    smallest = index;
    l = left(index);
    r = right(index);
 
    if (l < data.length && data[l] < data[smallest]) {
      smallest = l;
    }
    if (r < data.length && data[r] < data[smallest]) {
      smallest = r;
    }
    if (smallest == index) break;
 
    swap(index, smallest);
    index = smallest;
  }
}

Обе операции поднимают или опускают элемент не больше чем на высоту дерева, а высота бинарного дерева из n элементов равна O(log n). Поэтому push и pop работают за O(log n).

уровень 0: 1 узел
уровень 1: 2 узла
уровень 2: 4 узла
уровень 3: 8 узлов

Выше мы привели псевдокод, чтобы показать, как куча устроена внутри. На интервью и в реальном коде редко пишут кучу с нуля. В языках и библиотеках уже есть готовые реализации, и обычно достаточно знать, как создать кучу, добавить элементы и достать максимумы.

import "container/heap"
 
// MaxHeap - тип с методом Less(i, j int) bool { return h[i] > h[j] }
h := &MaxHeap{2, 7, 4, 1}
heap.Init(h)
 
heap.Push(h, 8)
 
max1 := heap.Pop(h).(int) // 8
max2 := heap.Pop(h).(int) // 7
 

Когда использовать

Куча часто подходит, если в задаче есть такие признаки:

1. Нужно постоянно получать минимум или максимум.

2. Нужно найти k наибольших или k наименьших элементов.

3. Нужно объединять несколько отсортированных структур.

4. Нужно обрабатывать поток данных.

5. Нужно на каждом шаге выбирать лучший доступный вариант.

Если нужно один раз отсортировать массив и пройтись по нему, обычная сортировка может быть проще. Куча становится особенно полезной, когда данные меняются во время работы алгоритма.

Пример

В задаче Вес последнего камня мы каждый раз берем два самых тяжелых камня. Если использовать сортировку после каждого удара, получится много лишней работы. Максимальная куча позволяет быстро достать два самых больших элемента, посчитать остаток и вернуть его обратно.

stones = [2, 7, 4, 1, 8, 1];

Первый шаг:

8 и 7 -> остается 1

Новый камень 1 снова попадает в кучу, и процесс продолжается.

Итоги

Куча нужна там, где важен приоритет, а не полный порядок. Она позволяет обрабатывать элементы постепенно и каждый раз быстро получать самый подходящий элемент.