Premium

Кучи

K-й наибольший в потоке

Поддерживать k наибольших элементов потока

Легкий

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

Описание (LeetCode - 703. Kth Largest Element in a Stream)

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

Вам нужно реализовать класс, который для заданного целого числа k поддерживает поток результатов тестов и непрерывно возвращает k-й наивысший результат теста после отправки нового результата.

Класс поддерживает две операции:

  • KthLargest(k, nums) - создает объект и инициализирует его числами из nums.
  • add(val) - добавляет новое число val в поток и возвращает k-й наибольший элемент среди всех полученных чисел.

Пример 1:

k = 3
nums = [4, 5, 8, 2]

Операции:

add(3)  -> 4
add(5)  -> 5
add(10) -> 5
add(9)  -> 8
add(4)  -> 8

Пояснение:

После add(3) поток содержит [4, 5, 8, 2, 3]. Три наибольших числа - 8, 5 и 4, поэтому третий наибольший элемент равен 4.

После add(5) три наибольших числа - 8, 5 и 5, поэтому ответ равен 5.

Ограничения:

  • 1 <= k <= 10⁴
  • 0 <= nums.length <= 10⁴
  • -10⁴ <= nums[i] <= 10⁴
  • -10⁴ <= val <= 10⁴
  • Будет выполнено не более 10⁴ вызовов add.

Решение

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

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

Будем поддерживать минимальную кучу размера не более k. В этой куче хранятся текущие k наибольших элементов потока.

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

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

Минимальная куча размера k

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

  1. Создаем минимальную кучу.
  2. Добавляем в нее все числа из nums с помощью метода add.
  3. В add(val) кладем val в кучу.
  4. Если размер кучи превышает k, удаляем минимум.
  5. Возвращаем вершину кучи.

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

k = 3
nums = [4, 5, 8, 2]
  1. После инициализации куча хранит три наибольших элемента: [4, 5, 8]. Минимум среди них равен 4, значит, текущий третий наибольший элемент - 4.
После инициализации
  1. Добавим 3. В куче временно окажутся [3, 4, 5, 8].
add(3): добавляем в кучу

Размер кучи превысил k, поэтому удаляем 3. В куче остаются [4, 5, 8], ответ равен 4.

add(3): удаляем минимум
  1. Добавим 5. В куче временно окажутся [4, 5, 5, 8].
add(5): добавляем в кучу

Удаляем 4. Теперь минимум среди трех лучших равен 5, поэтому ответ равен 5.

add(5): удаляем минимум

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

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

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

Добавление одного элемента в кучу размера k занимает O(log k).

Инициализация массивом из n элементов занимает O(n log k), а каждый последующий вызов add выполняется за O(log k).

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

Куча хранит не более k элементов, поэтому пространственная сложность равна O(k).

Код решения

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

import "github.com/emirpasic/gods/trees/binaryheap"
 
type KthLargest struct {
    k    int
    heap *binaryheap.Heap
}
 
func Constructor(k int, nums []int) KthLargest {
    kth := KthLargest{k: k, heap: binaryheap.NewWithIntComparator()}
    for _, num := range nums {
        kth.Add(num)
    }
    return kth
}
 
func (this *KthLargest) Add(val int) int {
    this.heap.Push(val)
    if this.heap.Size() > this.k {
        this.heap.Pop()
    }
    peek, _ := this.heap.Peek()
    return peek.(int)
}

Итоги

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

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