K-я порядковая статистика за линейное время

$k$-я порядковая статистика — это элемент, который окажется на $k$-м месте, если массив отсортировать по возрастанию.

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

Оказывается, можно за $O(n)$.

Идея

Возьмём случайный элемент $x$ и посчитаем, сколько в массиве элементов, не превосходящих его. Пусть их $cnt$.

Если $cnt \ge k$, то $k$-й по возрастанию элемент не больше $x$: маленьких элементов набралось достаточно, и искомый — среди них. Значит, искать надо только в левой половине.

Если $cnt < k$, то элементов, не больших $x$, слишком мало, и искомый лежит правее. Но номер придётся пересчитать: среди правой половины нам нужен уже не $k$-й, а $(k - cnt)$-й — первые $cnt$ элементов остались слева.

В обоих случаях половина массива отбрасывается насовсем.

Код

Разделение — то же самое partition, что и в быстрой сортировке.

// k-й по возрастанию (нумерация с единицы) в полуинтервале [l, r).
int kth_element(vector<int>& a, int l, int r, int k) {
    if (r - l == 1) return a[l];
    int pos = partition(a, l, r);
    int cnt = pos - l;                    // столько ушло в левую половину
    if (cnt >= k) return kth_element(a, l, pos, k);
    return kth_element(a, pos, r, k - cnt);
}

Базовый случай — один элемент: если отрезок сузился до него, он и есть ответ.

Почему это линия

Разница с быстрой сортировкой — в одной строке: там два рекурсивных вызова, здесь один. И этого достаточно, чтобы сменилась асимптотика.

Возьмём идеальный случай: каждый раз делим пополам. Тогда работа складывается так:

$$n + \frac{n}{2} + \frac{n}{4} + \frac{n}{8} + \ldots$$

Это сумма геометрической прогрессии, и она меньше $2n$ — сколько бы слагаемых ни было. Не $n \log n$, а именно $2n$: уровней по-прежнему логарифм, но каждый следующий вдвое дешевле предыдущего, и всё вместе сходится.

Сравните с сортировкой: там на каждом уровне работа снова равна $n$, потому что обрабатываются обе половины.

Уровень	Сортировка: обе половины	Поиск $k$-го: одна половина
1	$n$	$n$
2	$n$	$n/2$
3	$n$	$n/4$
…	…	…
Итого	$n \log n$	меньше $2n$

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

А без случайности можно?

Всё сказанное опирается на случайный выбор опорного: гарантии тут вероятностные. Существует и детерминированный алгоритм с честным $O(n)$ в худшем случае — «медиана медиан».

Идея: разбить массив на группы по пять элементов, в каждой найти медиану (это $O(1)$ на группу), затем рекурсивно найти медиану этих медиан и взять её опорным. Такой опорный гарантированно отсекает хотя бы $30\%$ массива, откуда и берётся линейность.

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

В C++ это уже написано

nth_element(a.begin(), a.begin() + k, a.end());
// теперь a[k] — тот элемент, который стоял бы там после сортировки

Функция ничего не возвращает, а переставляет массив: элемент с индексом k встаёт на своё место, слева от него оказываются все не большие, справа — все не меньшие. Внутри половин порядка нет, и он не нужен.

Работает она в среднем за линию — то есть заметно быстрее, чем sort, если вам действительно нужен один элемент, а не весь порядок.