Редакция для Медиана потока


Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
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 → позиция 0
  • k = 2 → позиция 0
  • k = 3 → позиция 1
  • k = 4 → позиция 1

То есть при чётном количестве элементов берётся "левая" из двух середин.

Наблюдение 2. Как должны делиться элементы

Пусть:

  • в left_heap лежит левая половина,
  • в right_heap лежит правая половина.

Тогда нужно поддерживать два условия:

  1. Размеры куч отличаются не более чем на 1, причём left_heap может быть либо равна right_heap, либо на 1 больше.
  2. Каждый элемент из left_heap не больше любого элемента из right_heap.

Тогда:

  • если элементов нечётное количество, в left_heap будет на один элемент больше;
  • если чётное, размеры будут равны;
  • в обоих случаях нижняя медиана — это максимум в left_heap.

Наблюдение 3. Почему нужны именно две кучи

Куча умеет:

  • быстро добавлять элемент;
  • быстро брать и удалять минимум или максимум.

Это позволяет после каждого нового числа за O(log n):

  • положить его в подходящую половину;
  • при необходимости перебалансировать кучи;
  • сразу получить медиану.

3. Алгоритм

Будем хранить:

  • left_heap — max-heap;
  • right_heap — min-heap.

Для каждого нового числа x:

  1. Если left_heap пустая или x <= top(left_heap), кладём x в left_heap. Иначе кладём x в right_heap.

  2. Восстанавливаем размеры:

    • если left_heap.size() < right_heap.size(), переносим минимум из right_heap в left_heap;
    • если left_heap.size() > right_heap.size() + 1, переносим максимум из left_heap в right_heap.
  3. Дополнительно проверяем порядок:

    • если обе кучи не пусты и top(left_heap) > top(right_heap), значит граница между половинами нарушена;
    • тогда меняем эти два верхних элемента местами.
  4. Печатаем top(left_heap) — это текущая нижняя медиана.

Особенность в Python

В Python стандартная heapq — это min-heap.
Чтобы сделать max-heap для левой кучи, храним числа с обратным знаком:

  • в куче лежит -x,
  • настоящий максимум равен -left[0].

4. Почему это работает

Докажем, что после каждого шага вершина left_heap действительно равна нижней медиане.

Инварианты

После обработки очередного элемента поддерживаются два свойства:

  1. left_heap содержит либо столько же элементов, сколько right_heap, либо на один больше.
  2. Любой элемент из 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])

Комментарии

Еще нет ни одного комментария.