Поиск k-го элемента в несортированном массиве

Поиск k-го элемента в несортированном массиве

Задача

Найти k-й по величине элемент в несортированном массиве за время лучше чем O(n log n).

Для примера find_kth([3, 2, 1, 5, 6, 4], 2) ожидается 5 — второй по величине элемент. Отсортированный массив: [1, 2, 3, 4, 5, 6], второй с конца — это 5.

Решение 1: Quickselect (O(n) в среднем)

QuickSelect — это адаптация алгоритма QuickSort. Вместо полной сортировки, мы используем partitioning для поиска позиции нужного элемента.

import random

def find_kth_quickselect(arr, k):
    """
    Находит k-й по величине элемент за O(n) в среднем.
    k=1 означает максимум, k=n означает минимум.
    """
    if not arr or k < 1 or k > len(arr):
        return None
    
    # Преобразуем k в индекс (k-й по величине = индекс n-k)
    target_index = len(arr) - k
    
    def quickselect(left, right, target):
        if left == right:
            return arr[left]
        
        # Выбираем случайный pivot для избежания худшего случая O(n²)
        pivot_index = random.randint(left, right)
        pivot_index = partition(left, right, pivot_index)
        
        if target == pivot_index:
            return arr[target]
        elif target < pivot_index:
            return quickselect(left, pivot_index - 1, target)
        else:
            return quickselect(pivot_index + 1, right, target)
    
    def partition(left, right, pivot_index):
        pivot_value = arr[pivot_index]
        # Перемещаем pivot в конец
        arr[pivot_index], arr[right] = arr[right], arr[pivot_index]
        store_index = left
        
        for i in range(left, right):
            if arr[i] < pivot_value:
                arr[store_index], arr[i] = arr[i], arr[store_index]
                store_index += 1
        
        # Перемещаем pivot на его финальную позицию
        arr[right], arr[store_index] = arr[store_index], arr[right]
        return store_index
    
    return quickselect(0, len(arr) - 1, target_index)

# Тестирование
print(find_kth_quickselect([3, 2, 1, 5, 6, 4], 2))  # 5
print(find_kth_quickselect([3, 2, 3, 1, 2, 4, 5, 5, 6], 4))  # 4

Сложность:

• Лучший/средний случай: O(n) (в среднем pivot делит массив поровну)
• Худший случай: O(n²) (если pivot всегда в краю), но редко с рандомизацией

Решение 2: Heap (O(n log k))

Используем min-heap для отслеживания k-наибольших элементов.

import heapq

def find_kth_heap(arr, k):
    """
    Находит k-й по величине элемент за O(n log k).
    Эффективно когда k << n.
    """
    if not arr or k < 1 or k > len(arr):
        return None
    
    # Используем min-heap размером k
    min_heap = arr[:k]
    heapq.heapify(min_heap)
    
    # Для каждого оставшегося элемента
    for i in range(k, len(arr)):
        if arr[i] > min_heap[0]:  # Если больше минимума в heap
            heapq.heapreplace(min_heap, arr[i])
    
    # Минимум в heap — это k-й по величине элемент
    return min_heap[0]

# Тестирование
print(find_kth_heap([3, 2, 1, 5, 6, 4], 2))  # 5
print(find_kth_heap([3, 2, 3, 1, 2, 4, 5, 5, 6], 4))  # 4

Сложность: O(n log k), что даёт O(n log n) в худшем случае, но очень эффективно при малых k.

Решение 3: Сортировка (O(n log n))

Просто, но не лучше чем требуемое:

def find_kth_sort(arr, k):
    """
    Самое простое решение: O(n log n).
    """
    if not arr or k < 1 or k > len(arr):
        return None
    
    arr.sort(reverse=True)
    return arr[k - 1]

# Тестирование
print(find_kth_sort([3, 2, 1, 5, 6, 4], 2))  # 5
print(find_kth_sort([3, 2, 3, 1, 2, 4, 5, 5, 6], 4))  # 4

Решение 4: QuickSelect без модификации исходного массива

def find_kth_safe(arr, k):
    """
    QuickSelect, не изменяющий исходный массив.
    """
    if not arr or k < 1 or k > len(arr):
        return None
    
    def quickselect(elements, left, right, target):
        if left == right:
            return elements[left]
        
        pivot_index = random.randint(left, right)
        pivot_index = partition(elements, left, right, pivot_index)
        
        if target == pivot_index:
            return elements[target]
        elif target < pivot_index:
            return quickselect(elements, left, pivot_index - 1, target)
        else:
            return quickselect(elements, pivot_index + 1, right, target)
    
    def partition(elements, left, right, pivot_index):
        pivot_value = elements[pivot_index]
        # Не изменяем исходный массив
        temp = elements[:]
        temp[pivot_index], temp[right] = temp[right], temp[pivot_index]
        store_index = left
        
        for i in range(left, right):
            if temp[i] < pivot_value:
                temp[store_index], temp[i] = temp[i], temp[store_index]
                store_index += 1
        
        temp[right], temp[store_index] = temp[store_index], temp[right]
        # Возвращаем оба: обновлённый массив и индекс
        elements[:] = temp
        return store_index
    
    target_index = len(arr) - k
    return quickselect(arr, 0, len(arr) - 1, target_index)

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

Сравнение подходов

Медиана как частный случай

Нахождение медианы — это поиск элемента на позиции n/2:

def find_median(arr):
    """Находит медиану массива за O(n)"""
    n = len(arr)
    # Для нечётного размера
    if n % 2 == 1:
        return find_kth_quickselect(arr, n // 2 + 1)
    # Для чётного размера
    else:
        mid1 = find_kth_quickselect(arr, n // 2)
        mid2 = find_kth_quickselect(arr, n // 2 + 1)
        return (mid1 + mid2) / 2

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

Заключение

QuickSelect — лучший выбор для задачи, так как:

1. Достигает O(n) в среднем случае
2. Лучше чем O(n log n) требование
3. С рандомизацией, практически никогда не O(n²)
4. Эффективнее Heap при больших k

Для интервью рекомендуется сначала предложить QuickSelect, затем обсудить Heap как альтернативу для малых k.