Find Median From Data Stream

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

Надо реализовать следующий интерфейс:

// получает новый элемент из потока
void addNum(int num)
// находит медиану среди «прошедших через addNum значений»
double findMedian()

Пример:

addNum(1)
addNum(2)
// среднее значение между [1,2] –> 1.5
findMedian() -> 1.5
addNum(3) 
// среднее значение между [1,2,3] –> 2
findMedian() -> 2

Задача на LeetCode.

Решение #

Очевидное брутфорсное решение — складывать числа в коллекцию, и, когда просят найти медиану, сортировать и находить центр. Однако, в условии не просто так сказано про поток: со временем данных в коллекции будет всё больше и больше, и сортировать её всё дольше и дольше.

По сути, нам и не нужны данные в отсортированном виде, просто у нас нет пока нет подходящей структуры данных, которая бы позволяла найти «центральный элемент».

Хорошая практика при решении задач — найти «бутылочное горлышко». В чем оно заключается сейчас? Когда приходит очередной число надо понять в какую половину оно относится, по определению медианы. Однако, нам не нужен строгий порядок между всеми числами.

Представим, что у мы каким-то образом поделили числа на две половины (спойлер: прямо так и просится слово «куча» сюда). Приходит новое число. Смотрим прищуренным взглядом на кучи и понимаем, это число явно идёт в большую или меньшую половины. Нам не надо пересортировать все числа друг относительно друга.

Есть специальная структура данных, которая так и называется heap (куча) — дерево, которое в корне хранит наибольшее или наименьшее значение среди всех узлов. Сложность вставки и удаления элементов: O(log N).

Здесь мы и избавляемся от «бутылочного горлышка» — log N вместо N * log N при сортировке. Отношения порядка между элементами нет, важно поставить куда надо только один элемент (минимальный или максимальный).

Это именно то, что нужно! Заводим две кучи и поддерживаем их размеры отличающимися не более чем на единицу.

Далее прямо по определению медианы. Либо размеры куч одинаковые, т.е. количество чисел чётное, и тогда надо взять вершины деревьев (за O(1)) и найти среднее арифметическое. Либо размеры куч отличаются на единицу, т.е. количество чисел нечётное, и тогда надо взять вершину бо́льшего дерева.

Написать Heap отличное упражнение, но на собеседовании стоит воспользоваться стандартной библиотекой. В случае с JavaScript традиционно можно сказать «представим, что реализован следующий интерфейс». Главное обсудить как это работает.

Моя реализация вышесказанного через очереди с приоритетом на C++.

class MedianFinder {
private:
    // делим все числа на две кучи: maxH и minH,
    // одна держит максимальный элемент в корне,
    // а другая соответственно — минимальный
    priority_queue<int> maxH;
    priority_queue<int, vector<int>, greater<int>> minH;
public:
    void addNum(int num) {
        // не теряя общности закинем
        // новое число в одну половину
        maxH.push(num);
        // взамен перекинем вершину
        // в другую половину
        minH.push(maxH.top());
        // не забываем удалить элемент,
        // т.к. по смыслу он
        // теперь в другой куче
        maxH.pop();

        // ни одна куча не должна расти быстрее другой,
        // т.к. мы кидаем элементы в minH нужно убедиться,
        // что maxH при этом не пустеет
        if (minH.size() > maxH.size()) {
            maxH.push(minH.top());
            minH.pop();
        }
    }
    
    double findMedian() {
        // чётное количество:
        // возьмём среднее арифметическое
        if (minH.size() == maxH.size()) {
            return (minH.top() + maxH.top()) / 2.0;
        }
        // нечётное количество:
        // возьмём вершину большего дерева
        // на данном этапе мы знаем,
        // что в левой должно быть ровно на 1 больше
        return maxH.top();
    }
};

Красота! По-моему, довольно компактное решение.

Сложность — логарифм по времени и линия по памяти. У задачи есть два любопытных развития:

Спойлер: bucket sort

PS. Обсудить можно в телеграм-чате любознательных программистов. Welcome! 🤗

Подписывайтесь на мой твитер или канал в телеграме, чтобы узнавать о новых разборах задач.