Метод двух куч для медианы потока: как поддерживать медиану при добавлении чисел?
Числа приходят по одному, после каждого нужно вывести текущую медиану всех пришедших. Сортировать каждый раз — O(n^2 log n), слишком долго. Слышал про «две кучи», но не понимаю, как они дают медиану.
2 ответа
Держим две кучи, делящие отсортированный набор пополам:
lo— max-heap на меньшую половину (вершина = максимум левой части).hi— min-heap на большую половину (вершина = минимум правой части).
Инвариант: все элементы lo <= всех элементов hi, и размеры различаются не больше чем на 1 (lo.size() = hi.size() или +1). Тогда медиана — это вершина lo (при нечётном числе) или среднее вершин обеих куч (при чётном).
priority_queue<int> lo; // max-heap
priority_queue<int, vector<int>, greater<int>> hi; // min-heap
void add(int x) {
if (lo.empty() || x <= lo.top()) lo.push(x);
else hi.push(x);
// балансировка размеров
if (lo.size() > hi.size() + 1) { hi.push(lo.top()); lo.pop(); }
else if (hi.size() > lo.size()) { lo.push(hi.top()); hi.pop(); }
}
double median() {
if (lo.size() > hi.size()) return lo.top();
return (lo.top() + (double)hi.top()) / 2.0;
}
Вставка O(log n), запрос медианы O(1). На n чисел — суммарно O(n log n), память O(n).
Логика: новый элемент кладём в ту половину, куда он по значению попадает (сравнив с вершиной lo), затем восстанавливаем баланс размеров, перекидывая одну вершину. Инвариант «lo <= hi» держится автоматически, потому что перекидываем именно граничные вершины.
Если в задаче есть ещё и удаление произвольного элемента (скользящая медиана окна), обычная priority_queue не умеет удалять из середины. Тогда либо multiset с итератором на медиану (двигаешь его на каждый add/remove), либо «ленивое удаление»: вторая пара куч-«отложенных», и при извлечении вершины проверяешь, не помечена ли она удалённой. Для чистого потока без удалений двух куч хватает и они быстрее. Грабля: priority_queue<int> по умолчанию — max-heap, для min-heap нужен greater<int> третьим параметром, легко перепутать и получить WA.