Редакция для Медиана потока
Submitting an official solution before solving the problem yourself is a bannable offence.
1. Идея
Нужно после каждого нового числа быстро находить текущую медиану.
Если каждый раз сортировать все уже полученные числа заново, получится слишком медленно: до O(n^2 log n) или O(n^2) в зависимости от способа.
Ключевая идея — хранить числа в двух кучах:
left_heap— все числа, которые относятся к "левой половине", то есть не больше медианы;right_heap— все числа, которые относятся к "правой половине", то есть больше медианы.
При этом:
- в
left_heapудобно быстро брать максимальный элемент; - в
right_heapудобно быстро брать минимальный элемент.
Так как по условию нужна нижняя медиана, ответом всегда будет максимум из левой половины. Значит, после каждого добавления достаточно вывести вершину left_heap.
2. Наблюдения
Наблюдение 1. Что такое нижняя медиана
Для k чисел нужно взять элемент на позиции (k - 1) / 2 в отсортированном массиве.
Примеры:
k = 1→ позиция0k = 2→ позиция0k = 3→ позиция1k = 4→ позиция1
То есть при чётном количестве элементов берётся "левая" из двух середин.
Наблюдение 2. Как должны делиться элементы
Пусть:
- в
left_heapлежит левая половина, - в
right_heapлежит правая половина.
Тогда нужно поддерживать два условия:
- Размеры куч отличаются не более чем на 1, причём
left_heapможет быть либо равнаright_heap, либо на 1 больше. - Каждый элемент из
left_heapне больше любого элемента изright_heap.
Тогда:
- если элементов нечётное количество, в
left_heapбудет на один элемент больше; - если чётное, размеры будут равны;
- в обоих случаях нижняя медиана — это максимум в
left_heap.
Наблюдение 3. Почему нужны именно две кучи
Куча умеет:
- быстро добавлять элемент;
- быстро брать и удалять минимум или максимум.
Это позволяет после каждого нового числа за O(log n):
- положить его в подходящую половину;
- при необходимости перебалансировать кучи;
- сразу получить медиану.
3. Алгоритм
Будем хранить:
left_heap— max-heap;right_heap— min-heap.
Для каждого нового числа x:
Если
left_heapпустая илиx <= top(left_heap), кладёмxвleft_heap. Иначе кладёмxвright_heap.Восстанавливаем размеры:
- если
left_heap.size() < right_heap.size(), переносим минимум изright_heapвleft_heap; - если
left_heap.size() > right_heap.size() + 1, переносим максимум изleft_heapвright_heap.
- если
Дополнительно проверяем порядок:
- если обе кучи не пусты и
top(left_heap) > top(right_heap), значит граница между половинами нарушена; - тогда меняем эти два верхних элемента местами.
- если обе кучи не пусты и
Печатаем
top(left_heap)— это текущая нижняя медиана.
Особенность в Python
В Python стандартная heapq — это min-heap.
Чтобы сделать max-heap для левой кучи, храним числа с обратным знаком:
- в куче лежит
-x, - настоящий максимум равен
-left[0].
4. Почему это работает
Докажем, что после каждого шага вершина left_heap действительно равна нижней медиане.
Инварианты
После обработки очередного элемента поддерживаются два свойства:
left_heapсодержит либо столько же элементов, сколькоright_heap, либо на один больше.- Любой элемент из
left_heapне больше любого элемента изright_heap.
Покажем, почему они сохраняются.
Добавление элемента
Когда приходит новое число x, мы сначала кладём его:
- в
left_heap, если оно не больше текущего максимума левой части; - иначе в
right_heap.
После этого размеры могут немного нарушиться, но не более чем на 1.
Балансировка размеров
Если правая куча стала больше, переносим её минимум в левую.
Если левая стала слишком большой, переносим её максимум в правую.
После этого выполняется условие по размерам:
left_heap.size() == right_heap.size()илиleft_heap.size() == right_heap.size() + 1.
Исправление порядка
После переносов может оказаться, что максимум слева больше минимума справа.
Это означает, что некоторые элементы оказались "не в своей половине".
Тогда достаточно обменять местами верхушки куч:
- самый большой слева уходит направо,
- самый маленький справа уходит налево.
После этого снова будет верно, что все элементы слева не больше всех элементов справа.
Почему ответ — это вершина left_heap
Рассмотрим общее число обработанных элементов k.
Если k нечётно
Тогда left_heap содержит на один элемент больше, чем right_heap.
Значит, в левой куче лежат первые (k + 1) / 2 элементов отсортированного порядка, а её максимум — это элемент на позиции (k - 1) / 2, то есть медиана.
Если k чётно
Тогда размеры куч равны.
В left_heap лежат первые k / 2 элементов, а её максимум — это элемент на позиции k / 2 - 1.
Это и есть нижняя медиана, потому что (k - 1) / 2 = k / 2 - 1 при чётном k.
Значит, после каждого шага top(left_heap) — правильный ответ.
5. Сложность
Для каждого из n чисел выполняется несколько операций с кучами, каждая из которых работает за O(log n).
Итого:
- время:
O(n log n) - память:
O(n)
Это подходит для ограничения n <= 2 * 10^5.
6. Код на C++17
#include <iostream>
#include <vector>
#include <queue>
using namespace std;
int main() {
int n;
cin >> n;
priority_queue<long long> left_heap;
priority_queue<long long, vector<long long>, greater<long long>> right_heap;
for (int i = 0; i < n; i++) {
long long x;
cin >> x;
if (left_heap.empty() || x <= left_heap.top()) {
left_heap.push(x);
} else {
right_heap.push(x);
}
if ((int)left_heap.size() < (int)right_heap.size()) {
left_heap.push(right_heap.top());
right_heap.pop();
} else if ((int)left_heap.size() > (int)right_heap.size() + 1) {
right_heap.push(left_heap.top());
left_heap.pop();
}
if (!left_heap.empty() && !right_heap.empty() && left_heap.top() > right_heap.top()) {
long long a = left_heap.top();
long long b = right_heap.top();
left_heap.pop();
right_heap.pop();
left_heap.push(b);
right_heap.push(a);
}
cout << left_heap.top() << '\n';
}
return 0;
}
7. Код на Python 3
import heapq
n = int(input())
a = list(map(int, input().split()))
left = []
right = []
for x in a:
if not left or x <= -left[0]:
heapq.heappush(left, -x)
else:
heapq.heappush(right, x)
if len(left) < len(right):
heapq.heappush(left, -heapq.heappop(right))
elif len(left) > len(right) + 1:
heapq.heappush(right, -heapq.heappop(left))
if left and right and (-left[0] > right[0]):
l = -heapq.heappop(left)
r = heapq.heappop(right)
heapq.heappush(left, -r)
heapq.heappush(right, l)
print(-left[0])
Комментарии