Median of Two Sorted Arrays: медиана двух массивов за O(log min(m, n))

Коротко

• Главная идея: медиана делит объединённый массив на две равные половины. Достаточно найти, сколько элементов левой половины взять из первого массива, — остальное возьмётся из второго.
• Бинарный поиск идёт по более короткому массиву, поэтому сложность зависит от $\min(m, n)$.
• Половина ошибок — в граничных случаях: пустой массив, разбиение у края, чётная суммарная длина.

Условие

Даны два отсортированных по неубыванию массива nums1 длины $m$ и nums2 длины $n$. Нужно вернуть медиану их объединения. В оригинальной формулировке LeetCode (задача 4, уровень Hard) требуется сложность $O(\log(m + n))$.

nums1 = [1, 3], nums2 = [2]       → 2.0
nums1 = [1, 2], nums2 = [3, 4]    → 2.5
nums1 = [],     nums2 = [1]       → 1.0

Напомним, что такое медиана. Если элементов нечётное число, это средний элемент отсортированной последовательности. Если чётное — среднее арифметическое двух средних.

Задача знаменита тем, что наивное решение пишется за пять минут, а оптимальное многие не могут довести до рабочего состояния и за час. Разница — в аккуратности с индексами.

Решение 1: слияние за O(m + n)

Самое прямолинейное: слить два отсортированных массива в один (как в шаге сортировки слиянием) и взять середину.

def find_median_merge(nums1, nums2):
    merged = []
    i = j = 0
    while i < len(nums1) and j < len(nums2):
        if nums1[i] <= nums2[j]:
            merged.append(nums1[i])
            i += 1
        else:
            merged.append(nums2[j])
            j += 1
    merged.extend(nums1[i:])
    merged.extend(nums2[j:])

    n = len(merged)
    if n % 2 == 1:
        return float(merged[n // 2])
    return (merged[n // 2 - 1] + merged[n // 2]) / 2

Время $O(m + n)$, память $O(m + n)$. Вариант sorted(nums1 + nums2) короче, но даёт $O((m+n)\log(m+n))$ и выбрасывает информацию о том, что массивы уже отсортированы. На собеседовании его лучше упомянуть только как отправную точку.

Решение 2: два указателя без лишней памяти

Сливать целиком не нужно: медиана лежит на позиции около $(m+n)/2$. Пройдём двумя указателями ровно до неё, запоминая текущий и предыдущий элемент.

def find_median_two_pointers(nums1, nums2):
    m, n = len(nums1), len(nums2)
    total = m + n
    i = j = 0
    prev = cur = 0
    for _ in range(total // 2 + 1):
        prev = cur
        if i < m and (j >= n or nums1[i] <= nums2[j]):
            cur = nums1[i]
            i += 1
        else:
            cur = nums2[j]
            j += 1
    if total % 2 == 1:
        return float(cur)
    return (prev + cur) / 2

Память $O(1)$, время всё ещё линейное. Условие j >= n or ... закрывает случай, когда один массив закончился раньше. Это хороший промежуточный ответ: интервьюер видит, что вы понимаете структуру задачи, и обычно спрашивает: «А быстрее можно?»

Решение 3: k-я порядковая статистика

Обобщим задачу: найти $k$-й по счёту элемент объединения двух отсортированных массивов. Медиана — это $k = \lceil (m+n)/2 \rceil$ (и ещё следующий элемент при чётной длине).

Идея: сравним элементы на позиции $k/2$ в обоих массивах. Если a[k/2 - 1] <= b[k/2 - 1], то первые $k/2$ элементов массива a точно не могут быть $k$-м элементом: перед любым из них стоит не больше $k/2 - 1$ элементов из a и не больше $k/2 - 1$ из b, итого меньше $k - 1$. Их можно отбросить и искать $(k - k/2)$-й элемент в остатке. Каждый шаг уменьшает $k$ вдвое.

def kth(a, b, k):
    """k-й по счёту (с 1) элемент объединения двух отсортированных массивов."""
    i = j = 0
    while True:
        if i == len(a):
            return b[j + k - 1]
        if j == len(b):
            return a[i + k - 1]
        if k == 1:
            return min(a[i], b[j])
        half = k // 2
        ni = min(i + half, len(a)) - 1
        nj = min(j + half, len(b)) - 1
        if a[ni] <= b[nj]:
            k -= ni - i + 1
            i = ni + 1
        else:
            k -= nj - j + 1
            j = nj + 1


def find_median_kth(nums1, nums2):
    total = len(nums1) + len(nums2)
    if total % 2 == 1:
        return float(kth(nums1, nums2, total // 2 + 1))
    left = kth(nums1, nums2, total // 2)
    right = kth(nums1, nums2, total // 2 + 1)
    return (left + right) / 2

Тонкое место — min(i + half, len(a)): если в массиве осталось меньше $k/2$ элементов, берём сколько есть. Отбрасываем ровно столько, сколько реально прошли, поэтому k уменьшается на ni - i + 1, а не на half.

Сложность — $O(\log k) = O(\log(m+n))$, то есть формально требование задачи выполнено. Тот же приём «отбросить заведомо лишнюю часть» лежит в основе quickselect, только там отбрасывание вероятностное.

Решение 4: бинарный поиск по разбиению

Это решение, которое ждут на собеседовании. Оно быстрее и, если понять идею, пишется без рекурсии и без двух вызовов.

Интуиция

Разрежем nums1 после $i$ элементов, а nums2 — после $j$ элементов:

nums1:  left1_part (i штук)  |  right1_part
nums2:  left2_part (j штук)  |  right2_part

Левая половина объединения — это left1_part + left2_part. Чтобы она действительно была «левой половиной отсортированного объединения», нужны два условия:

1. Размер: $i + j = \lfloor (m + n + 1) / 2 \rfloor$. При нечётной сумме левая половина на один элемент больше — тогда медиана лежит в ней.
2. Порядок: каждый элемент слева не больше каждого элемента справа. Внутри одного массива это и так верно, поэтому достаточно проверить перекрёстно:

Из первого условия $j$ однозначно определяется через $i$. Значит, перебирать нужно только $i$ от $0$ до $m$ — и делать это бинарным поиском:

• если nums1[i-1] > nums2[j] — из nums1 взяли слишком много, двигаем i влево;
• если nums2[j-1] > nums1[i] — взяли слишком мало, двигаем i вправо.

Монотонность здесь очевидна: чем больше $i$, тем больше nums1[i-1] и тем меньше nums2[j]. Это та же схема, что и в бинарном поиске по ответу: ищем границу, где условие начинает выполняться.

Когда разбиение найдено:

• нечётная длина: медиана = $\max(\text{left1}, \text{left2})$;
• чётная: медиана = $\big(\max(\text{left1}, \text{left2}) + \min(\text{right1}, \text{right2})\big) / 2$.

Код

def find_median_sorted_arrays(nums1, nums2):
    if len(nums1) > len(nums2):
        nums1, nums2 = nums2, nums1
    m, n = len(nums1), len(nums2)
    half = (m + n + 1) // 2
    lo, hi = 0, m

    while lo <= hi:
        i = (lo + hi) // 2  # сколько берём из nums1
        j = half - i        # сколько берём из nums2

        left1 = nums1[i - 1] if i > 0 else float('-inf')
        right1 = nums1[i] if i < m else float('inf')
        left2 = nums2[j - 1] if j > 0 else float('-inf')
        right2 = nums2[j] if j < n else float('inf')

        if left1 <= right2 and left2 <= right1:
            if (m + n) % 2 == 1:
                return float(max(left1, left2))
            return (max(left1, left2) + min(right1, right2)) / 2
        if left1 > right2:
            hi = i - 1
        else:
            lo = i + 1

    raise ValueError('массивы не отсортированы')

Бесконечности решают проблему краёв: если из массива ничего не взяли слева ($i = 0$), левого соседа нет и он не должен мешать проверке — подставляем $-\infty$. Аналогично справа $+\infty$.

Почему именно по короткому массиву

Обмен nums1 и nums2 в начале — не косметика. Если nums1 длиннее, то при некоторых $i$ получится $j = \text{half} - i < 0$ или $j > n$, и индексация nums2[j - 1] сломается (в Python — молча возьмёт элемент с конца). Когда $m \le n$, для любого $i \in [0, m]$ выполняется $0 \le j \le n$:

Бонус — сложность становится $O(\log \min(m, n))$, это лучше, чем требует условие.

Разбор на примере

nums1 = [1, 3, 8, 9, 15], nums2 = [7, 11, 18, 19, 21, 25]. Всего 11 элементов, half = 6, nums1 короче — ищем $i \in [0, 5]$.

Левая половина: [1, 3, 8, 9] + [7, 11], правая: [15] + [18, 19, 21, 25]. Длина нечётная, медиана = $\max(9, 11) = 11$. Проверка слиянием: [1, 3, 7, 8, 9, 11, 15, 18, 19, 21, 25], шестой элемент — 11.

То же решение на JavaScript (у вас могут попросить другой язык):

function findMedianSortedArrays(nums1, nums2) {
  if (nums1.length > nums2.length) [nums1, nums2] = [nums2, nums1]
  const m = nums1.length
  const n = nums2.length
  const half = (m + n + 1) >> 1
  let lo = 0
  let hi = m
  while (lo <= hi) {
    const i = (lo + hi) >> 1
    const j = half - i
    const left1 = i > 0 ? nums1[i - 1] : -Infinity
    const right1 = i < m ? nums1[i] : Infinity
    const left2 = j > 0 ? nums2[j - 1] : -Infinity
    const right2 = j < n ? nums2[j] : Infinity
    if (left1 <= right2 && left2 <= right1) {
      if ((m + n) % 2 === 1) return Math.max(left1, left2)
      return (Math.max(left1, left2) + Math.min(right1, right2)) / 2
    }
    if (left1 > right2) hi = i - 1
    else lo = i + 1
  }
  throw new Error('arrays are not sorted')
}

Граничные случаи

Прежде чем сказать «готово», прогоните решение по этому списку вслух. Интервьюеры часто спрашивают именно про них.

Ещё два вопроса, которые стоит уточнить у интервьюера в начале:

• Тип результата. На LeetCode ответ — double. В Java и C++ деление (a + b) / 2 на целых числах обрежет дробную часть: пишите (a + b) / 2.0.
• Переполнение. В языках с фиксированной разрядностью a + b может переполнить int, если значения близки к границе. Безопаснее a / 2.0 + b / 2.0 или приведение к long. В Python проблемы нет. Заглушки $\pm\infty$ в Java обычно заменяют на Integer.MIN_VALUE / Integer.MAX_VALUE: в итоговую формулу они никогда не попадают, потому что при найденном корректном разбиении хотя бы один из двух кандидатов — настоящий элемент.

Сложность

Для $m = n = 10^6$ разница между линейным и логарифмическим решением — два миллиона операций против примерно двадцати итераций.

Типичные ошибки

1. Бинарный поиск по длинному массиву. $j$ уходит в минус, Python молча берёт nums2[-1], тесты иногда проходят, иногда нет. Всегда меняйте массивы местами.
2. hi = m - 1 вместо hi = m. Разбиение $i = m$ («взять весь nums1 влево») — допустимый вариант. Если отрезать его, сломается пример [1, 2], [3, 4].
3. Неправильный half. С (m + n) // 2 при нечётной длине медиана окажется справа, и формулу max(left1, left2) придётся менять на min(right1, right2). Оба варианта рабочие, но смешивать их нельзя. Выберите один и проговорите его.
4. Строгие сравнения. left1 < right2 вместо <= ломает решение на дубликатах: равные элементы по обе стороны разреза — нормальная ситуация.
5. Заглушки 0 или -1 вместо бесконечностей. Работают только на положительных данных.
6. Целочисленное деление в Java/C++ — уже упоминали выше, но это самая частая причина Wrong Answer у тех, кто пишет не на Python.

Как рассказывать на собеседовании

Хорошая последовательность — та же, что в статье:

1. Уточнить: массивы отсортированы, может ли один быть пустым, какой тип ответа.
2. Назвать слияние за $O(m+n)$ и сразу оговорить, что оно не использует логарифм из условия.
3. Сформулировать ключевое наблюдение: медиана = разрез объединения на две равные половины; разрез задаётся одним числом $i$.
4. Объяснить монотонность, из-за которой работает бинарный поиск, и почему он идёт по короткому массиву.
5. Написать код, проговорить бесконечности на краях и пройти пример из таблицы выше.

Если вы застряли на индексах, не молчите: проговорите инвариант (что означают i, j, half) — это лучший способ выбраться и одновременно показать ход мыслей. Об этом подробнее в статье про рассуждения вслух.

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

• k-й элемент двух отсортированных массивов — прямое обобщение, решается функцией kth выше.
• Медиана потока данных — когда элементы приходят по одному, а не лежат в массивах. Разбор с двумя кучами — в статье Find Median from Data Stream.
• Поиск в повёрнутом отсортированном массиве — ещё одна задача, где бинарный поиск идёт не по значению, а по структурному свойству: Search in Rotated Sorted Array.
• Основы бинарного поиска и шаблоны границ — в статье бинарный поиск.