K самых частых элементов

Постановка задачи

nums = [1, 1, 1, 2, 2, 3]
k = 2

[1, 2]

nums = [1]
k = 1

[1]

Решение

В этой задаче нужно решить две подзадачи:

1. Посчитать частоту каждого элемента (с помощью хеш-таблицы).
2. Выбрать k элементов с наибольшей частотой (с помощью кучи).

Идея решения:

Сначала считаем частоты всех чисел с помощью хеш-таблицы. Затем поддерживаем минимальную кучу размера k, где каждый элемент хранит пару (число, частота), упорядоченную по частоте.

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

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

Основные шаги:

1. Считаем частоту каждого числа в хеш-таблице frequency.
2. Создаем минимальную кучу, упорядоченную по частоте.
3. Проходим по парам (число, частота) из хеш-таблицы:
   • добавляем пару в кучу;
   • если размер кучи больше k, удаляем элемент с наименьшей частотой.
4. Извлекаем оставшиеся числа из кучи и возвращаем результат.

Рассмотрим пример:

nums = [1, 1, 1, 2, 2, 3]
k = 2

Частоты:

1 -> 3
2 -> 2
3 -> 1

1. Добавляем 1 с частотой 3. Куча: [(1, 3)].

2. Добавляем 2 с частотой 2. Куча: [(2, 2), (1, 3)].

3. Добавляем 3 с частотой 1. Размер стал 3, а нам нужно хранить не более 2 элементов. Удаляем вершину с минимальной частотой - пару (3, 1).

В куче остались 1 и 2, это и есть ответ.

Оценка сложности

Временная сложность

Подсчет частот занимает O(n). Пусть m - количество различных элементов в массиве (m <= n). Для каждого из них мы выполняем добавление в кучу и, возможно, удаление за O(log k).

Итоговая временная сложность: O(n + m log k). В худшем случае m = n, поэтому сложность составляет O(n log k).

Пространственная сложность

Хеш-таблица хранит до m различных чисел, а куча хранит не более k элементов.

Итоговая дополнительная память: O(m + k), что в худшем случае равно O(n).

Код решения

Приведем код решения.

import "github.com/emirpasic/gods/trees/binaryheap"
 
type Pair struct {
    num  int
    freq int
}
 
func topKFrequent(nums []int, k int) []int {
    frequency := map[int]int{}
    for _, num := range nums {
        frequency[num]++
    }
 
    heap := binaryheap.NewWith(func(a, b interface{}) int {
        return a.(Pair).freq - b.(Pair).freq
    })
 
    for num, freq := range frequency {
        heap.Push(Pair{num: num, freq: freq})
        if heap.Size() > k {
            heap.Pop()
        }
    }
 
    result := []int{}
    for !heap.Empty() {
        val, _ := heap.Pop()
        result = append(result, val.(Pair).num)
    }
    return result
}

Итоги

Задача Top K Frequent Elements объединяет хеш-таблицу для подсчета частот и кучу для выбора лучших кандидатов.

1. Суть алгоритма: сначала считаем частоты в хеш-таблице, затем поддерживаем минимальную кучу размера k по этим частотам.
2. Эффективность:
   • Временная сложность: O(n + m log k), где m - количество различных элементов.
   • Пространственная сложность: O(m + k).
3. Граничные случаи:
   • k = 1: в куче останется один самый частый элемент.
   • Все элементы уникальны: ответом будут любые k элементов, так как частота каждого равна 1.
4. Оптимальное решение: мы не сортируем все элементы по частоте. Куча размера k позволяет сразу отбрасывать кандидатов с меньшей частотой.