K-й наименьший элемент в двух отсортированных массивах

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

Даны два отсортированных массива arr1 и arr2 и целое число k. Требуется вернуть k-й наименьший элемент из объединённого отсортированного массива. Эта задача — классика собеседований и прямой наследник задачи о медиане двух отсортированных массивов.

Брутфорс-решение (слияние массивов)

На собеседовании можно начать с простого объяснения: так как оба массива отсортированы, мы можем слить их как в сортировке слиянием. Во время слияния считаем элементы, и как только доходим до k-го, возвращаем его.

Сложность: O(N+M) по времени и O(N+M) по памяти.

def kth_element_brute(arr1, arr2, k):
    merged = []
    i = j = 0
    while i < len(arr1) and j < len(arr2):
        if arr1[i] < arr2[j]:
            merged.append(arr1[i])
            i += 1
        else:
            merged.append(arr2[j])
            j += 1
    # добавляем остатки
    merged.extend(arr1[i:])
    merged.extend(arr2[j:])
    return merged[k-1]

Оптимальный подход: бинарный поиск по разбиению

Заметим, что для задачи «Медиана двух отсортированных массивов» мы разбивали массивы так, чтобы слева оказалась половина всех элементов. Для k-го элемента нужно, чтобы слева оказалось ровно k элементов. Идея разбиения остаётся той же.

Ключевое наблюдение

Пусть arr1 = [2,3,6,7,9], arr2 = [1,4,8,10], k = 5. Нам нужно, чтобы слева было ровно 5 элементов. Рассмотрим разбиение:

arr1: 2 3 6 | 7 9
arr2:   1 4 | 8 10
Левая сторона: 1 2 3 4 6 — ровно 5 элементов.
Ответ: max(6,4) = 6

Условия корректного разбиения: left1 <= right2 и left2 <= right1. Если разбиение валидно, то ответ — max(left1, left2).

Оптимальное решение на Python

Бинарный поиск ведём по меньшему массиву. Параметры: cut1 — сколько элементов берём из первого массива, cut2 = k - cut1 — из второго. Границы бинарного поиска: low = max(0, k - n2), high = min(k, n1).

def kth_element(arr1, arr2, k):
    n1, n2 = len(arr1), len(arr2)
    # гарантируем, что n1 <= n2
    if n1 > n2:
        return kth_element(arr2, arr1, k)

    low = max(0, k - n2)
    high = min(k, n1)

    while low <= high:
        cut1 = (low + high) // 2
        cut2 = k - cut1

        left1 = arr1[cut1 - 1] if cut1 > 0 else float('-inf')
        left2 = arr2[cut2 - 1] if cut2 > 0 else float('-inf')
        right1 = arr1[cut1] if cut1 < n1 else float('inf')
        right2 = arr2[cut2] if cut2 < n2 else float('inf')

        if left1 <= right2 and left2 <= right1:
            # валидное разбиение
            return max(left1, left2)
        elif left1 > right2:
            # нужно уменьшить cut1
            high = cut1 - 1
        else:
            # нужно увеличить cut1
            low = cut1 + 1

    return -1  # не должно случиться

Сухой прогон

Для примера выше: low = max(0,5-4)=1, high = min(5,5)=5. Первая итерация: cut1=3, cut2=2. left1=6, left2=4, right1=7, right2=8. Условия: 6<=8 и 4<=7 — верно, ответ max(6,4)=6.

Почему это работает?

k-й элемент — это наибольший среди первых k элементов. Когда разбиение валидно, левая сторона содержит ровно k элементов, и максимальный из них — искомый.

Анализ сложности

• Время: O(log(min(N, M))) — бинарный поиск по меньшему массиву.
• Память: O(1) — используем только переменные.

Паттерн для собеседований

Если видите: два отсортированных массива + k-й наименьший элемент + ожидается логарифмическое решение → думайте «бинарный поиск по разбиению». Запомните: медиана — это частный случай с k = (N+M)/2.

Похожие задачи

• Медиана двух отсортированных массивов
• K-й наименьший элемент в отсортированной матрице
• Поиск в сдвинутом массиве

Практический вывод

Прямо сейчас откройте LeetCode или аналогичный ресурс и решите задачу «Kth Element of Two Sorted Arrays» (или «Median of Two Sorted Arrays» для разминки). Используйте технику разбиения — она пригодится и в других задачах.