Исправляем Quickselect: корректное разбиение при дубликатах

Почему Quickselect с разбиением Ломуто тормозит при дубликатах и K в середине, и как исправить: рандомизация равных и трёхпутевое разбиение без деградации.

Quickselect — базовый прием для поиска K‑го по величине элемента в массиве с ожидаемой линейной сложностью. Но даже он может упираться в лимит времени на вполне разумных данных. Если разбиение выполняется с фиксированным опорным элементом на конце и равные значения обрабатываются «в лоб», то большой массив с множеством дубликатов и K около середины способен загнать алгоритм в квадратичный режим. Понимание причины и способов её нейтрализовать — ключ к надежной реализации выбора ранга.

Постановка задачи

Рассмотрим прямолинейный Quickselect, который перемешивает входной массив, делает разбиение по схеме Ломуто с последним элементом в роли опорного и рекурсивно углубляется в ту часть, где находится искомый индекс. Поиск K‑го наибольшего сводится к индексу K‑го наименьшего через целевую позицию n − k при нулевой индексации.

import random

class Ranker(object):
    def kth_largest(self, arr, kpos):
        m = len(arr)
        random.shuffle(arr)
        return self._pick(arr, 0, m - 1, m - kpos)

    def _pick(self, arr, lo, hi, idx):
        p = self._split(arr, lo, hi)
        if idx < p:
            return self._pick(arr, lo, p - 1, idx)
        elif idx > p:
            return self._pick(arr, p + 1, hi, idx)
        else:
            return arr[p]

    def _split(self, arr, lo, hi):
        p = hi
        i = lo
        for j in range(lo, hi):
            if arr[j] <= arr[p]:
                arr[i], arr[j] = arr[j], arr[i]
                i += 1
        arr[i], arr[p] = arr[p], arr[i]
        return i

Где именно всё ломается

Замедление проявляется на входах с низкой энтропией — например, когда массив почти целиком состоит из одного и того же значения, а целевой индекс расположен около середины. В приведённой выше схеме Ломуто условие разбиения использует «меньше-или-равно». Если опорное значение встречается много раз, цикл переносит каждый равный элемент в левую часть и в итоге ставит сам опорный в самый конец секции. После этого рекурсия отбрасывает за шаг всего один элемент. Такой сценарий повторяется до достижения цели, и суммарное время вырождается в квадратичное. Перемешивание не спасает, потому что проблема не только в выборе опоры, а в обработке равных значений при разбиении.

Как исправить разбиение для равных значений

Чтобы разорвать этот худший сценарий, в разбиении нельзя сливать все элементы, равные опоре, на одну сторону. Практичный приём: переставлять элементы строго по условию «меньше», а для значений, равных опоре, выполнять перестановку лишь в половине случаев. Такое случайное распределение равных по обе стороны границы не даёт отрезку сжиматься по одному элементу за шаг и устраняет тайм-ауты на больших массивах с множеством дубликатов.

import random

class Ranker(object):
    def kth_largest(self, arr, kpos):
        m = len(arr)
        random.shuffle(arr)
        return self._pick(arr, 0, m - 1, m - kpos)

    def _pick(self, arr, lo, hi, idx):
        p = self._split(arr, lo, hi)
        if idx < p:
            return self._pick(arr, lo, p - 1, idx)
        elif idx > p:
            return self._pick(arr, p + 1, hi, idx)
        else:
            return arr[p]

    def _split(self, arr, lo, hi):
        p = hi
        i = lo
        for j in range(lo, hi):
            if arr[j] < arr[p] or (arr[j] == arr[p] and random.getrandbits(1)):
                arr[i], arr[j] = arr[j], arr[i]
                i += 1
        arr[i], arr[p] = arr[p], arr[i]
        return i

Небольшое изменение сохраняет корректность и не меняет общую структуру Quickselect, но устраняет патологию, вызванную скоплениями одинаковых значений. Решение остаётся лёгким: без дополнительных проходов и памяти.

Почему это важно

В реальных данных дубликатов обычно много. Если разбиение уносит все равные в одну сторону, Quickselect будет снова и снова «снимать по одному элементу» и уйдёт в квадратичное время. Такой скрытый обрыв производительности легко не заметить при разработке — он проявляется лишь на крупных данных или на некоторых наборах тестов. Поэтому корректная обработка равных — не микрооптимизация, а страховка от худших сценариев.

Существуют и другие стратегии разбиения, которые хорошо работают при обилии дублей. Схема Хоара обычно ведёт себя лучше по мере роста числа повторов, а трёхпутевое разбиение (вариант «голландского флага») — распространённый способ эффективно отделять области «меньше», «равно» и «больше». Во всех подходах идея одна: не позволять равным значениям перекосить рекурсию.

Итоги

Если Quickselect вылетает по времени на больших массивах при крупном K, внимательно проверьте, как разбиение обрабатывает равные. Избегайте «воронки меньше-или-равно» в схеме Ломуто: рандомизируйте размещение равных или перейдите на метод, который балансирует их изначально. Тестируйте на массивах, где доминирует одно повторяющееся число и цель находится около середины, — так вы убедитесь, что реализация не деградирует. С такой защитой Quickselect остаётся надёжным инструментом для выбора K‑го по величине элемента.