Skip to main content

Top K

Top-K pattern — общий приём для задач «найди k лучших»

Что это за задачи

В интервью часто встречается формулировка вида:

«Дан набор объектов и число k. Верните k “лучших” объектов, где “лучше” определено условием задачи (частота появления, наибольший/наименьший вес, максимальный рейтинг и т. д.).»

Классический «топ-K» появляется в задачах про частоты слов, самые большие/маленькие значения, лучшие результаты и т. п.

Наивное решение

  1. Посчитать нужную метрику для каждого объекта.
  2. Отсортировать весь набор по этой метрике.
  3. Взять первые k элементов.

Сложность — O(n log n), где n — размер входа.

Улучшение с помощью кучи

Если k ≪ n, сортировать всё невыгодно. Приём «движущегося мини-куча» сокращает логарифм до log k:

  1. Создаём min-heap (или max-heap, если нужно искать k наименьших).

  2. Итерируем вход:

    • помещаем объект в кучу;
    • как только размер > k, делаем pop() — удаляется «худший» элемент.
  3. После прохода кучу составляют ровно k «лучших» объектов.

Сложность становится O(n log k) по времени и O(k) по памяти, что выигрывает у сортировки, когда k заметно меньше n.

Ключевая идея: ограничиваем размер приоритетной очереди k, чтобы все операции над ней стоили log k, а не log n.

Когда применять

  • k значительно меньше n;
  • нужно вернуть топ без строгого порядка (если важна сортировка результата, его можно отсортировать в конце — это k log k, обычно несущественно);
  • в условии нет удаления произвольных элементов (при удалениях сложнее, лучше multiset/treap).

Такой шаблон лежит в основе решений для задач «Top K Frequent Elements», «K Closest Points to Origin», «Kth Largest Element in an Array» и множества вариаций.

Top K Frequent Elements

Условие задачи

Дан массив целых чисел nums и число k. Нужно вернуть k элементов, которые встречаются чаще всего. Гарантируется, что ответ единственен (частоты различаются или наборы не пересекаются).


Ключевая идея

  1. Посчитать частоту каждого числа хеш-картой freq.

  2. Поддерживать min-кучу размера ≤ k, где ключ — частота, значение — число.

    • Пока обходим freq, кладём в кучу пару (frequency, value).
    • Как только размер стал k + 1, делаем pop() — уходит элемент с минимальной частотой.
  3. После обхода в куче останутся ровно k самых частотных чисел.

Поскольку размер очереди ограничен k, все операции с ней стоят O(log k), а не O(log n).


Почему не хватает одной только сортировки

Сортировка словаря по убыванию частоты даёт O(n log n). С кучей получаем O(n log k) (потому что log k < log n), что особенно ощутимо, когда k во много раз меньше n.


Алгоритм

  1. freq = Map() → для каждого num увеличиваем счётчик.

  2. Создать MinPriorityQueue с приоритетом pair => pair[0] (частота).

  3. Для каждой [val, f] из freq:

    1. enqueue([f, val])
    2. если queue.size() > kdequeue()
  4. Извлечь из очереди k значений и вернуть их в любом порядке.


Реализация на JavaScript

npm i @datastructures-js/priority-queue
/**
* @param {number[]} nums
* @param {number} k
* @return {number[]}
*/
var topKFrequent = function (nums, k) {
const map = new Map();
const minHeap = new MinPriorityQueue((keyVal) => keyVal[1]);
const ans = [];

for (let i = 0; i < nums.length; i++) {
map.set(nums[i], (map.get(nums[i]) || 0) + 1);
}

for (const [key, val] of map) {
minHeap.enqueue([key, val]);
if (minHeap.size() > k) {
minHeap.dequeue();
}
}

while (minHeap.size() > 0) {
ans.push(minHeap.dequeue()[0]);
}

return ans;
};

Анализ сложности

ЭтапВремяПамять
Подсчёт частотO(n)O(n)
Работа c очередью (n вставок)O(n log k)O(k)
ИтогоO(n log k)O(n + k)

n = nums.length. Хеш-карта хранит все частоты, очередь — лишь k элементов, что обеспечивает улучшение по времени по сравнению с полной сортировкой.