Медиана за окно
Метеоролог сглаживает ряд суточных температур скользящей медианой: для каждого окна из w дней подряд берёт средний по величине день. В отличие от среднего арифметического, медиану не сдвигает одна аномальная оттепель или ошибка датчика. Если в окне чётное число дней, берётся нижняя медиана — меньший из двух средних.
Два теста — настоящая среднесуточная температура Алматы с 1 января 1941 года по 31 декабря 2025-го в десятых долях градуса, 31 046 дней, с окном в год и в месяц.
Вход
В первой строке n и w (1 ≤ w ≤ n ≤ 100 000). Во второй — n целых чисел a1, …, an, |ai| ≤ 10⁹.
Выход
n − w + 1 чисел, по одному в строке: для каждого i от 1 до n − w + 1 — элемент, который стоит на месте ⌊(w + 1)/2⌋ (считая с единицы), если числа ai, …, ai+w−1 упорядочить по возрастанию.
Примеры
Пример 1
8 3 5 -3 8 8 1 -7 4 9
5 8 8 1 1 4
Пример 2
7 4 2 2 7 2 7 7 1
2 2 7 2
Ограничения
- Время · Python
- 4000 мс
- Время · Go
- 1000 мс
- Память
- 256 МБ
- Проверка
- 11 тестов
Разбор
Открывайте по ступеням: каждая следующая подсказывает больше. Можно сразу к решению — но задача, решённая с одной подсказкой, запоминается лучше.
1Подсказка 1
Без удалений задача известная: медиану потока держат две кучи — в нижней меньшая половина чисел с наибольшим наверху, в верхней большая половина с наименьшим наверху. Медиана — вершина нижней. Что мешает здесь?
2Подсказка 2
День, выпавший из окна, может лежать в глубине любой из куч. Не ищите его там — удалите лениво: пометьте и выбросьте, когда он окажется наверху. Но тогда размеры куч врут. Что нужно считать отдельно?
3Подсказка 3
Одинаковые температуры мешают понять, в какой куче лежит выпавший день. Сделайте все ключи разными.
4Идея
Ключ дня i — пара (значение, i): одинаковых ключей больше нет, и все живые ключи нижней кучи меньше всех живых ключей верхней. Тогда выпавший день лежит в нижней куче ровно тогда, когда его ключ не больше её живой вершины. Считаем живые ключи в каждой куче отдельно; удаление — это пометка и уменьшение счётчика своей кучи. Помеченные ключи выбрасываем, только когда они на вершине. После каждого шага выравниваем кучи так, чтобы в нижней было ⌊(w + 1)/2⌋ живых ключей; её вершина и есть нижняя медиана.
5Разбор
Пару (значение, i) удобно упаковать в одно число: (ai + 10⁹)·n + i. Порядок таких чисел — порядок пар, номер дня восстанавливается как остаток от деления на n, значение — как частное минус 10⁹. Нижняя куча хранит ключи с минусом, чтобы heapq и container/heap держали наверху наибольший.
Шаг для дня i: положить его ключ в нижнюю кучу, если он меньше её живой вершины, иначе в верхнюю. Если день i − w выпал из окна, узнать, в какой он куче (сравнить с живой вершиной нижней — до того, как пометить его), уменьшить счётчик этой кучи и пометить день. Потом перекладывать живую вершину из кучи в кучу, пока в нижней не станет ровно ⌊(k + 1)/2⌋ живых ключей, где k — сколько дней сейчас в окне. Функция «живая вершина» сначала выбрасывает помеченные ключи сверху, потом возвращает вершину.
Есть и решение без куч: держать окно отсортированным списком, выпавший день удалять, а новый вставлять на место двоичным поиском (bisect.insort). Каждый шаг стоит O(w), но это один сдвиг памяти, и, как показывает замер в статье, такие сдвиги очень быстрые: при n до 100 000 решение укладывается в ограничения, на пятом тесте оно впятеро медленнее куч. На ряде в десятки миллионов значений с окном в миллион сдвиги проиграли бы кучам в сотни раз.
Мусор в кучах ограничен: каждый день попадает в кучу один раз и один раз из неё выбрасывается. Даже если помеченные ключи надолго застрянут в глубине, куча не больше n.
5+Решение
6Сложность
На каждый день — несколько операций с кучами, каждая O(log n); выбрасывание помеченных ключей за весь проход стоит O(n log n). Итого O(n log n) времени и O(n) памяти.
7Типичные ошибки
- Сортировать каждое окно. O(n · w log w): в пятом тесте пятьдесят тысяч окон по пятьдесят тысяч дней.
- Пометить, а потом решать, где лежал. Если выпавший день был живой вершиной нижней кучи, выбрасывание помеченных уберёт его, и сравнение с новой вершиной отправит его счёт не в ту кучу.
- Сравнивать по значению без номера дня. Можно и так, но тогда при одинаковых температурах «лежит в нижней куче» уже не следует из сравнения, и счётчики нужно вести иначе — через разницу размеров. Ключи с номером дня снимают вопрос целиком.
- Верхняя медиана вместо нижней. При чётном w во втором примере нужна меньшая из двух средних — в нижней куче должно быть ⌊(w + 1)/2⌋ ключей, а не ⌊w/2⌋.