Find Median from Data Stream: медиана потока на двух кучах

Коротко

• Делим все числа на две половины: меньшую храним в max-куче, большую — в min-куче. Медиана — на вершинах куч.
• Инвариант: размеры куч отличаются не больше чем на один, и любой элемент слева не больше любого справа.
• Для скользящей медианы добавляется удаление — его делают «лениво», через словарь отложенных удалений.

Условие

Нужно реализовать класс MedianFinder (LeetCode 295, Hard):

• addNum(num) — добавить число из потока;
• findMedian() — вернуть медиану всех добавленных чисел.

addNum(1), addNum(2) → findMedian() = 1.5
addNum(3)            → findMedian() = 2.0

Ключевое слово — поток. Числа приходят по одному, и медиану могут спросить в любой момент. Если массив дан целиком, хватит quickselect или сортировки. Здесь же важна стоимость каждой операции на протяжении всей жизни объекта.

Наивные решения

Пересортировывать при каждом запросе. findMedian за $O(n \log n)$. При $5 \cdot 10^4$ вызовах (ограничение в задаче) это слишком медленно.

Держать отсортированный список. Вставка через бинарный поиск:

from bisect import insort


class MedianFinderSorted:
    def __init__(self):
        self.data = []

    def addNum(self, num):
        insort(self.data, num)  # поиск O(log n), но сдвиг O(n)

    def findMedian(self):
        n = len(self.data)
        if n % 2 == 1:
            return float(self.data[n // 2])
        return (self.data[n // 2 - 1] + self.data[n // 2]) / 2

Место находится за $O(\log n)$, но вставка в середину массива сдвигает хвост — $O(n)$. На практике insort в Python работает быстро благодаря memmove, и на LeetCode такое решение проходит. Но на собеседовании ответ «у меня $O(n)$ на вставку» — повод для следующего вопроса.

Заметим, что для медианы не нужен полностью отсортированный массив. Нужно знать только наибольший элемент меньшей половины и наименьший элемент большей половины. «Быстро дать максимум» и «быстро дать минимум» — это ровно то, что умеет двоичная куча.

Две кучи

Идея

       low (max-куча)          high (min-куча)
   [1, 3, 5]  вершина = 5  |  вершина = 8  [8, 9, 12]
            меньшая половина  |  большая половина

Поддерживаем два инварианта:

1. Порядок: max(low) <= min(high).
2. Размеры: len(low) == len(high) или len(low) == len(high) + 1.

Тогда медиана:

• нечётное количество — вершина low;
• чётное — среднее вершин low и high.

Балансировка

Как добавить число и не сломать инварианты? Самый короткий и надёжный приём — «прогнать через обе кучи»:

1. Кладём число в low.
2. Достаём максимум из low и перекладываем в high. Теперь порядок гарантирован: в high попал наибольший элемент левой части.
3. Если high стал больше low, перекладываем минимум high обратно в low.

В Python heapq реализует только min-кучу, поэтому max-кучу эмулируют, храня числа с обратным знаком.

import heapq


class MedianFinder:
    def __init__(self):
        self.low = []   # max-куча (храним с минусом): меньшая половина
        self.high = []  # min-куча: большая половина

    def addNum(self, num: int) -> None:
        heapq.heappush(self.low, -num)
        heapq.heappush(self.high, -heapq.heappop(self.low))
        if len(self.high) > len(self.low):
            heapq.heappush(self.low, -heapq.heappop(self.high))

    def findMedian(self) -> float:
        if len(self.low) > len(self.high):
            return float(-self.low[0])
        return (-self.low[0] + self.high[0]) / 2

Каждая вставка — не больше трёх операций с кучей, то есть $O(\log n)$. Медиана — $O(1)$, это просто чтение вершин.

Альтернатива с меньшим числом операций: сравнить num с вершиной low, положить в нужную кучу, а потом выровнять размеры одним перекладыванием. Работает так же асимптотически, но в ней больше ветвлений, где легко ошибиться. На собеседовании короткий вариант выигрывает.

Разбор на примере

Поток 5, 15, 1, 3:

На последнем шаге: 3 попадает в low → [1, 3, 5], максимум 5 уходит в high → [5, 15]. Размеры 2 и 2, балансировать не нужно, медиана $(3 + 5) / 2 = 4$.

Тот же класс на Java

В Java max-куча делается компаратором, без трюка со знаком:

class MedianFinder {
    private final PriorityQueue<Integer> low =
        new PriorityQueue<>(Collections.reverseOrder());
    private final PriorityQueue<Integer> high = new PriorityQueue<>();

    public void addNum(int num) {
        low.offer(num);
        high.offer(low.poll());
        if (high.size() > low.size()) low.offer(high.poll());
    }

    public double findMedian() {
        if (low.size() > high.size()) return low.peek();
        return ((long) low.peek() + high.peek()) / 2.0;
    }
}

Обратите внимание на (long) и 2.0: без них сумма двух больших int переполнится, а деление станет целочисленным.

Follow-up 1: все числа от 0 до 100

Классический дополнительный вопрос к этой задаче: «Что если все числа из потока лежат в диапазоне $[0, 100]$?»

Тогда кучи не нужны. Заводим массив счётчиков на 101 ячейку — это сортировка подсчётом, только «онлайн». Медиану ищем, накапливая счётчики, пока не дойдём до нужной позиции.

class MedianFinder100:
    def __init__(self):
        self.count = [0] * 101
        self.n = 0

    def addNum(self, num):
        self.count[num] += 1
        self.n += 1

    def _kth(self, k):  # k-й по счёту, с 1
        acc = 0
        for value, c in enumerate(self.count):
            acc += c
            if acc >= k:
                return value

    def findMedian(self):
        if self.n % 2 == 1:
            return float(self._kth(self.n // 2 + 1))
        return (self._kth(self.n // 2) + self._kth(self.n // 2 + 1)) / 2

addNum — $O(1)$, findMedian — $O(101)$, то есть константа, и памяти тоже константа, сколько бы чисел ни пришло. Если запросов медианы очень много, можно хранить текущую медиану и сдвигать указатель при вставке, но обычно достаточно назвать решение со счётчиками.

Follow-up 2: 99% чисел в диапазоне 0..100

Второй вариант того же вопроса: «А если 99% чисел в $[0, 100]$, но бывают выбросы?»

Ответ: счётчики для диапазона плюс два отдельных хранилища для выбросов — числа меньше 0 и больше 100 (например, по куче или отсортированному списку на каждую сторону, и ещё просто их количества). Медиана при таком распределении почти наверняка лежит внутри $[0, 100]$: чтобы она вышла за край, выбросов с одной стороны должно быть больше половины потока. Поэтому:

• знаем below — сколько чисел меньше 0;
• ищем позицию медианы $k$, вычитаем below и идём по счётчикам;
• если $k \le$ below или $k >$ below + в_диапазоне — медиана среди выбросов, и тогда смотрим в их хранилище.

Интервьюеру здесь важно не точное количество строк кода, а то, что вы заметили: асимметрия данных позволяет отдать основной поток простой структуре, а редкие случаи обработать медленнее.

Скользящая медиана (Sliding Window Median)

Родственная задача (LeetCode 480, Hard): дан массив nums и размер окна k. Нужно вернуть медиану каждого окна длины k.

nums = [1, 3, -1, -3, 5, 3, 6, 7], k = 3
→ [1.0, -1.0, -1.0, 3.0, 5.0, 6.0]

Теперь элементы не только приходят, но и уходят из окна. А куча не умеет быстро удалять произвольный элемент: найти его — $O(k)$.

Ленивое удаление

Решение — не удалять сразу. Записываем уходящий элемент в словарь delayed и вычёркиваем его, только когда он окажется на вершине кучи. При этом нужно отдельно хранить «живые» размеры куч — физическая длина списка теперь включает мусор.

import heapq
from collections import defaultdict


def median_sliding_window(nums, k):
    low, high = [], []          # max-куча (с минусом) и min-куча
    delayed = defaultdict(int)  # число -> сколько раз ждёт удаления
    low_size = high_size = 0    # «живые» элементы в каждой куче

    def prune(heap, sign):
        # снимаем с вершины элементы, помеченные на удаление
        while heap and delayed[sign * heap[0]]:
            delayed[sign * heap[0]] -= 1
            heapq.heappop(heap)

    def rebalance():
        nonlocal low_size, high_size
        if low_size > high_size + 1:
            heapq.heappush(high, -heapq.heappop(low))
            low_size -= 1
            high_size += 1
            prune(low, -1)
        elif low_size < high_size:
            heapq.heappush(low, -heapq.heappop(high))
            low_size += 1
            high_size -= 1
            prune(high, 1)

    def add(num):
        nonlocal low_size, high_size
        if not low or num <= -low[0]:
            heapq.heappush(low, -num)
            low_size += 1
        else:
            heapq.heappush(high, num)
            high_size += 1
        rebalance()

    def remove(num):
        nonlocal low_size, high_size
        delayed[num] += 1
        if num <= -low[0]:
            low_size -= 1
            if num == -low[0]:
                prune(low, -1)
        else:
            high_size -= 1
            if high and num == high[0]:
                prune(high, 1)
        rebalance()

    def median():
        if k % 2 == 1:
            return float(-low[0])
        return (-low[0] + high[0]) / 2

    for i in range(k):
        add(nums[i])
    result = [median()]
    for i in range(k, len(nums)):
        add(nums[i])
        remove(nums[i - k])
        result.append(median())
    return result

Почему это корректно:

• Элемент, помеченный в delayed, лежит в той куче, куда указывает сравнение с вершиной low: инвариант порядка гарантирует, что всё <= -low[0] находится слева.
• После каждой операции вершины обеих куч — живые элементы. Это обеспечивают вызовы prune после удаления с вершины и после перекладывания. Поэтому median() читает правильные значения.
• Помеченные элементы в глубине куч не мешают: в расчёт они не попадают, а размеры мы считаем отдельно.

Сложность — $O(n \log n)$ по времени: каждый элемент добавляется один раз и вычёркивается один раз. Память — $O(n)$ в худшем случае, потому что мусор может копиться в глубине кучи.

Вариант попроще: отсортированное окно

Если интервьюер не настаивает на кучах, окно можно держать отсортированным списком:

from bisect import bisect_left, insort


def median_sliding_window_sorted(nums, k):
    window = sorted(nums[:k])
    result = []
    for i in range(k, len(nums) + 1):
        if k % 2 == 1:
            result.append(float(window[k // 2]))
        else:
            result.append((window[k // 2 - 1] + window[k // 2]) / 2)
        if i == len(nums):
            break
        window.pop(bisect_left(window, nums[i - k]))
        insort(window, nums[i])
    return result

Это $O(n \cdot k)$ из-за сдвигов, но код короткий и почти не оставляет места для ошибок. В Python с библиотекой sortedcontainers (SortedList) то же решение работает за $O(n \log k)$, однако на собеседовании её может не оказаться под рукой. В C++ аналогичную роль играет multiset с итератором на медиану. Ещё варианты — дерево отрезков или дерево Фенвика по сжатым значениям с поиском $k$-го элемента.

Сама техника сдвига окна разобрана в статье про скользящее окно.

Что уточнить у интервьюера

Задача кажется полностью определённой, но несколько вопросов в начале экономят время и показывают зрелость:

• Соотношение вызовов. Если findMedian вызывают редко, а addNum — миллионы раз, можно копить числа в буфере и считать медиану quickselect'ом по требованию. Если наоборот — две кучи с $O(1)$ на чтение подходят идеально.
• Диапазон значений. Это прямой мост к follow-up про 0..100: маленький диапазон меняет всё решение.
• Нужно ли удаление. Если да — это уже скользящая медиана или задача на сбалансированное дерево, кучи без ленивого удаления не справятся.
• Объём данных. Обе кучи хранят все числа, память $O(n)$. Если поток бесконечный и памяти не хватает, точную медиану не получить; тогда говорят о приближённых методах — гистограммах по корзинам или скетчах квантилей вроде t-digest. Достаточно назвать направление, детали на алгоритмическом интервью обычно не спрашивают.
• Тип ответа. На LeetCode ожидается вещественное число; при чётном количестве медиана может быть дробной даже для целых входных данных.

Типичные ошибки

1. Забыли про знак в heapq. Положили num в low без минуса — получили вторую min-кучу, и медиана считается неправильно, но иногда случайно совпадает на маленьких тестах.
2. Нарушили только один инвариант. Выравнивают размеры, но не следят за порядком (или наоборот). Пример: поток 1, 2, 0 при правиле «класть в кучу поменьше, при равенстве — в low» без сравнения с вершинами даст low = [1, 0], high = [2] — повезло; а 3, 1, 2 уже даст неверную медиану. Приём «прогнать через обе кучи» закрывает оба инварианта сразу.
3. Целочисленное деление и переполнение в Java/C++ при (a + b) / 2.
4. В скользящей медиане — длина списка вместо живого размера. len(low) включает помеченный мусор, балансировка по нему ломается.
5. Не вызвали prune после перекладывания. Перекладывание может обнажить помеченный элемент на вершине, и следующая медиана прочитает удалённое значение.
6. Отвечают на follow-up «кучами поменьше». Для диапазона 0..100 ожидается переход к счётчикам, а не оптимизация куч.

Где встречается

Find Median from Data Stream — одна из самых популярных Hard-задач на кучи: её часто включают в подборки задач для собеседований в крупные компании, и её стоит знать наизусть вместе с follow-up. Сам приём «две кучи держат границу» повторяется в других задачах:

• IPO (LeetCode 502) и задачи планирования, где одна куча хранит доступные варианты, а другая — ожидающие.
• Медиана двух отсортированных массивов — та же идея разреза на половины, но для статичных данных и через бинарный поиск: Median of Two Sorted Arrays.
• Top K элементов — одна куча фиксированного размера вместо двух: Top K Frequent Elements.
• Устройство очередей с приоритетом и их реализация — в статье приоритетная очередь.

На собеседовании рассказывайте в таком порядке: отсортированный список и его $O(n)$ на вставку → наблюдение «нужны только две средние точки» → две кучи с двумя инвариантами → пример на четырёх числах → follow-up со счётчиками. Это укладывается в 25–30 минут и оставляет время на скользящее окно, если интервьюер решит его спросить.