Найти медиану среди элементов потока данных

 Источник: https://leetcode.com/problems/find-median-from-data-stream/

Медиана - это значение, которое находится в середине отсортированного списка целых чисел. Если размер списка четный, то медиана - это среднее арифметическое двух элементов из середины списка.

• Например, для arr = [2,3,4], медиана =3.
• Дляarr = [2,3], медиана = (2 + 3) / 2 = 2.5.

Заимплементировать класс MedianFinder такой, что:

• MedianFinder() iинициализирует объект MedianFinder.
• void addNum(int num) добавляет целое число num из потока данных в структуру данных.
• double findMedian() возвращает медиану полученных данных на текущий момент. 

Пример 1:

Дано
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]
Результат
[null, null, null, 1.5, null, 2.0]

Пояснение
MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1);    // arr = [1]
medianFinder.addNum(2);    // arr = [1, 2]
medianFinder.findMedian(); // return 1.5 (i.e., (1 + 2) / 2)
medianFinder.addNum(3);    // arr[1, 2, 3]
medianFinder.findMedian(); // return 2.0

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

• -105 <= num <= 105
• В построенной структуре данных должен быть хотя бы один элемент перед вызовом функцииfindMedian.
• Может бысть сделано максимально 5 * 104 вызовов функций addNum и findMedian.

 

Follow up:

• If all integer numbers from the stream are in the range [0, 100], how would you optimize your solution?
• If 99% of all integer numbers from the stream are in the range [0, 100], how would you optimize your solution?

Идея решения

Что если при получении потока данных, мы просто каким-то образом обновляем срединный элемент (медиану), а сортировка остальных элементов потока данных нас не интересует?

Структура днных, которая удовлетворяет этому пожеланию - это КУЧА (heap).

Добавление элемента в кучу занимает алгоритмичсеское время. А доступ к максимальному (или минимальному - зависит от кучи) элемента занимает констаное время.

Мы можем поддерживать две кучи:

• max-heap хранит наименьшие элементы (половину полученного потока);
• min-heap хранит вторую половину элементов потока, но те, которые наибольшие.

Тогда:

1. Все элементы в max-heap меньше или равны корневого элемента (назовем его $x$).
2. Все элементы в min-heap больше или равны корневого элемента (назовем его $y$).

Если в потоке данных четное число элементов, то мы берем корневые элементы обеих куч и вычисляем их среднее арифметическое. Если количество элементов нечетное, то берем корневой элемент max-кучи.

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

Имплементация

// max-куча с наименьши элементами, в корневом элементе находится максимальный
private PriorityQueue<Integer> small = new PriorityQueue<>(Collections.reverseOrder());
// min-куча с наибольшими элементами, в корневом - минимальный
private PriorityQueue<Integer> large = new PriorityQueue<>();
// в каждый момент времени устанавливаем флаг, если количество элементов четно (true)
// или нечетно (false)
// изначально 0 элементов - значит, четно
private boolean even = true;

public double findMedian() {
    if (even) // если четно, то среднее арифметическое корней обеих куч
        return (small.peek() + large.peek()) / 2.0;
    else // если нечетно, то корень кучи с наименьшими элементами
        return small.peek();
}

public void addNum(int num) {
    if (even) { // если уже прибыло четное количество элементов, то есть к учи равны
// новый элемент добавляем в кучу с максимальными элементами, она перестроится
// и в ее корне окажется минимальный элемент
        large.offer(num);
// этот минимальный элемент перенесем в кучу с минимальными элементами
// обе кучи перестроятся, 
// но в куче с минимальными элементами будет на 1 элемент больше
        small.offer(large.poll());
    } else { // нечетное количество
// в куче с минимальными элементами k+1 элементов, с максимальными k элементов
// добавляем ещё один элемент в кучу с минимальными элементами - их k+2
        small.offer(num);
// переносим наибольший элемент среди минимальных в кучу с макисмальными элементами
// теперь в обеих кучах по k+1 элементов
        large.offer(small.poll());
    }
// меняем флаг четности на противоположный
    even = !even;
}