Quickselect: поиск k-го по величине элемента за O(n)

Коротко

Задача

Найти k-й по величине элемент массива. Частный случай — медиана (k = n/2).

Очевидное решение: отсортировать и взять нужный индекс.

def kth_smallest_naive(nums, k):
    return sorted(nums)[k - 1]

Работает, но делает лишнюю работу: мы упорядочиваем весь массив, хотя нужен один элемент.

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

Идея

Выбираем опорный элемент и разбиваем массив: слева меньшие, справа большие. После разбиения опорный элемент стоит на своей окончательной позиции.

Дальше сравниваем эту позицию с k:

• совпала — ответ найден;
• позиция больше k — искомый элемент левее, рекурсия в левую часть;
• позиция меньше k — рекурсия в правую часть.

Ключевое отличие от быстрой сортировки: вторую половину мы не трогаем вообще.

Реализация

import random


def quickselect(nums: list[int], k: int) -> int:
    """k-й наименьший элемент, k начинается с 1."""
    if not 1 <= k <= len(nums):
        raise ValueError('k вне диапазона')

    arr = nums[:]           # не портим вход
    left, right = 0, len(arr) - 1
    target = k - 1          # переводим в индекс

    while True:
        if left == right:
            return arr[left]

        pivot_index = partition(arr, left, right)

        if pivot_index == target:
            return arr[pivot_index]
        elif pivot_index > target:
            right = pivot_index - 1
        else:
            left = pivot_index + 1


def partition(arr: list[int], left: int, right: int) -> int:
    """Разбиение Ломуто со случайным опорным элементом."""
    # случайный выбор защищает от вырожденных данных
    rand = random.randint(left, right)
    arr[rand], arr[right] = arr[right], arr[rand]

    pivot = arr[right]
    i = left

    for j in range(left, right):
        if arr[j] < pivot:
            arr[i], arr[j] = arr[j], arr[i]
            i += 1

    arr[i], arr[right] = arr[right], arr[i]
    return i


nums = [3, 2, 1, 5, 6, 4]
print(quickselect(nums, 1))   # 1 — минимум
print(quickselect(nums, 3))   # 3
print(quickselect(nums, 6))   # 6 — максимум

# медиана
print(quickselect(nums, len(nums) // 2 + 1))   # 4

Реализация итеративная — рекурсия здесь хвостовая, и цикл экономит стек.

Почему в среднем O(n)

Каждое разбиение стоит O(n), но обрабатываемая часть в среднем уменьшается вдвое:

n + n/2 + n/4 + n/8 + ... = 2n

Сумма геометрической прогрессии даёт O(n), а не O(n log n). Логарифм не появляется именно потому, что мы спускаемся в одну ветвь, а не в обе.

Худший случай O(n²) — если опорный элемент каждый раз оказывается минимальным или максимальным. На отсортированном массиве с выбором последнего элемента в качестве опорного это происходит гарантированно.

Отсюда случайный выбор опорного элемента в коде выше: он делает плохой сценарий крайне маловероятным и убирает зависимость от порядка входных данных.

Медиана медиан: гарантированный O(n)

Существует алгоритм выбора опорного элемента, гарантирующий линейное время в худшем случае.

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

def median_of_medians(arr: list[int], k: int) -> int:
    if len(arr) <= 5:
        return sorted(arr)[k]

    # медианы групп по 5
    medians = [
        sorted(arr[i:i + 5])[len(arr[i:i + 5]) // 2]
        for i in range(0, len(arr), 5)
    ]
    pivot = median_of_medians(medians, len(medians) // 2)

    lows = [x for x in arr if x < pivot]
    highs = [x for x in arr if x > pivot]
    equals = len(arr) - len(lows) - len(highs)

    if k < len(lows):
        return median_of_medians(lows, k)
    elif k < len(lows) + equals:
        return pivot
    else:
        return median_of_medians(highs, k - len(lows) - equals)


print(median_of_medians([3, 2, 1, 5, 6, 4], 2))   # 3

На практике этот алгоритм медленнее обычного quickselect со случайным опорным элементом: константа заметно больше. Его ценность теоретическая — доказательство, что задача решается за линейное время в худшем случае. В реальном коде используют случайный выбор или гибрид (introselect), который переключается на медиану медиан при подозрении на деградацию.

Задача о k наибольших элементах

Частая формулировка. Здесь есть выбор между тремя подходами.

import heapq
import random

nums = [random.randint(0, 1000) for _ in range(100000)]
k = 10

# 1. сортировка: O(n log n)
top_sorted = sorted(nums, reverse=True)[:k]

# 2. куча: O(n log k) — лучше при малом k
top_heap = heapq.nlargest(k, nums)

# 3. quickselect: O(n), но результат не отсортирован
threshold = quickselect(nums, len(nums) - k + 1)
top_qs = [x for x in nums if x >= threshold][:k]

Правило выбора:

• k мало (десятки при миллионах элементов) — куча, O(n log k);
• k сравнимо с n — quickselect, O(n);
• нужен отсортированный результат — сортировка или досортировка k элементов после quickselect.

Где применяется

Медиана и перцентили — расчёт p95 и p99 задержек в мониторинге без сортировки всех замеров.

numpy.partition — реализация именно этого алгоритма.

Внутри nth_element в стандартной библиотеке C++.

Фильтрация выбросов — найти границу верхних 1% значений.

Задачи вида «k ближайших точек» — quickselect по расстоянию.

Частые ошибки

Фиксированный опорный элемент. Выбор первого или последнего даёт O(n²) на отсортированных данных — а это очень частый вход в реальных задачах.

Путаница k-го наибольшего и k-го наименьшего. Для наибольшего нужен индекс n - k. Ошибка на границе легко проходит тесты на симметричных данных.

Модификация входного массива. Алгоритм переставляет элементы. Если вход нужен вызывающему коду — копируйте.

Индексация с единицы или с нуля. Определитесь и придерживайтесь: в коде выше k с единицы, а target переводится в индекс.

Дубликаты при разбиении Ломуто. На массиве из одинаковых элементов классическое разбиение вырождается. Помогает трёхстороннее разбиение (по принципу «голландского флага»).

Что запомнить

• Quickselect ищет k-й элемент за O(n) в среднем, спускаясь только в одну половину.
• Линейность получается из суммы n + n/2 + n/4 + … = 2n.
• Опорный элемент обязательно выбирайте случайно — иначе O(n²) на отсортированных данных.
• Медиана медиан даёт гарантированный O(n), но на практике медленнее.
• При маленьком k куча (heapq.nlargest) обычно выгоднее.

Пройди собеседование в топ-компанию

Платформа для подготовки

Решай алгоритмические задачи как профи

✓ Популярные алгоритмы✓ Разбор решений✓ AI помощь

Начать сейчас→