Поиск медианы двух отсортированных массивов

Поиск медианы двух отсортированных массивов

Описание задачи

Это классическая задача на бинарный поиск. Требуется O(log(m+n)) время, поэтому наивный подход (объединить и найти медиану) не подойдёт.

Медиана:

• Если длина нечётная: средний элемент
• Если чётная: среднее двух средних элементов

Подход 1: Бинарный поиск (оптимальный)

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

def findMedianSortedArrays(nums1: list[int], nums2: list[int]) -> float:
    # Убедимся, что nums1 меньший массив (для оптимизации)
    if len(nums1) > len(nums2):
        nums1, nums2 = nums2, nums1
    
    m, n = len(nums1), len(nums2)
    low, high = 0, m
    
    while low <= high:
        partition1 = (low + high) // 2
        partition2 = (m + n + 1) // 2 - partition1
        
        # Получаем граничные значения
        left1 = float('-inf') if partition1 == 0 else nums1[partition1 - 1]
        right1 = float('inf') if partition1 == m else nums1[partition1]
        
        left2 = float('-inf') if partition2 == 0 else nums2[partition2 - 1]
        right2 = float('inf') if partition2 == n else nums2[partition2]
        
        # Проверяем, правильное ли разделение
        if left1 <= right2 and left2 <= right1:
            # Нашли правильное разделение
            total_length = m + n
            
            if total_length % 2 == 0:
                # Чётная длина: среднее двух средних
                return (max(left1, left2) + min(right1, right2)) / 2
            else:
                # Нечётная длина: больший из левых
                return max(left1, left2)
        
        elif left1 > right2:
            # Слишком много элементов из nums1
            high = partition1 - 1
        else:
            # Слишком мало элементов из nums1
            low = partition1 + 1
    
    return -1  # Не должно произойти

# Примеры
print(findMedianSortedArrays([1, 3], [2]))           # 2.0
print(findMedianSortedArrays([1, 2], [3, 4]))        # 2.5
print(findMedianSortedArrays([], [1]))               # 1.0
print(findMedianSortedArrays([], [1, 2]))            # 1.5
print(findMedianSortedArrays([0, 0], [0, 0]))        # 0.0
print(findMedianSortedArrays([2], []))               # 2.0

Логика:

1. Убеждаемся, что nums1 ≤ nums2 по длине (оптимизация)

2. Используем бинарный поиск для нахождения позиции разделения в nums1

3. Вычисляем соответствующую позицию в nums2:

   partition2 = (m + n + 1) // 2 - partition1
   

   Формула гарантирует, что слева от разделения будет ⌈(m+n)/2⌉ элементов

4. Получаем граничные элементы:

   • left1, right1 — элементы слева и справа в nums1
   • left2, right2 — элементы слева и справа в nums2

5. Проверяем условие: left1 <= right2 AND left2 <= right1

   • Если верно: нашли правильное разделение
   • Если left1 > right2: сдвигаемся влево
   • Если left1 <= right2: сдвигаемся вправо

6. Вычисляем медиану:

   • Чётная длина: (max(left1, left2) + min(right1, right2)) / 2
   • Нечётная длина: max(left1, left2)

Сложность:

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

Подход 2: Двойной указатель (понимание)

Менее эффективно, но более понятно:

def findMedianSortedArrays_v2(nums1: list[int], nums2: list[int]) -> float:
    merged = []
    i, j = 0, 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 == 0:
        return (merged[n // 2 - 1] + merged[n // 2]) / 2
    else:
        return float(merged[n // 2])

print(findMedianSortedArrays_v2([1, 3], [2]))       # 2.0
print(findMedianSortedArrays_v2([1, 2], [3, 4]))    # 2.5

Проблема: O(m + n) по времени и пространству (не подходит к требованиям)

Подход 3: Оптимизация двойного указателя

Без полного объединения массивов:

def findMedianSortedArrays_v3(nums1: list[int], nums2: list[int]) -> float:
    m, n = len(nums1), len(nums2)
    total = m + n
    
    # Нам нужны два средних элемента (или один для нечётной длины)
    mid1_pos = (total - 1) // 2
    mid2_pos = total // 2
    
    i, j = 0, 0
    current = None
    prev = None
    
    # Заходим в цикл только до позиции второй медианы
    for pos in range(mid2_pos + 1):
        prev = current
        
        if i < m and (j >= n or nums1[i] <= nums2[j]):
            current = nums1[i]
            i += 1
        else:
            current = nums2[j]
            j += 1
    
    if total % 2 == 0:
        return (prev + current) / 2
    else:
        return float(current)

print(findMedianSortedArrays_v3([1, 3], [2]))       # 2.0
print(findMedianSortedArrays_v3([1, 2], [3, 4]))    # 2.5

Улучшение: O((m+n)/2) в среднем, но всё ещё O(m+n) в худшем

Тестирование

import unittest

class TestFindMedianSortedArrays(unittest.TestCase):
    def setUp(self):
        self.solution = findMedianSortedArrays
    
    def test_basic(self):
        self.assertEqual(self.solution([1, 3], [2]), 2.0)
        self.assertEqual(self.solution([1, 2], [3, 4]), 2.5)
    
    def test_one_empty(self):
        self.assertEqual(self.solution([], [1]), 1.0)
        self.assertEqual(self.solution([1], []), 1.0)
        self.assertEqual(self.solution([], [1, 2]), 1.5)
    
    def test_both_empty(self):
        # Зависит от реализации — может быть исключение
        pass
    
    def test_different_lengths(self):
        self.assertEqual(self.solution([0, 0], [0, 0]), 0.0)
        self.assertEqual(self.solution([2], []), 2.0)
        self.assertEqual(self.solution([1, 3, 5], [2, 4]), 3.0)
    
    def test_large_arrays(self):
        nums1 = list(range(0, 1000, 2))  # [0, 2, 4, ..., 998]
        nums2 = list(range(1, 1000, 2))  # [1, 3, 5, ..., 999]
        # Медиана всех чисел от 0 до 999 — это 499.5
        self.assertAlmostEqual(self.solution(nums1, nums2), 499.5)

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

Ключевые моменты

1. Формула partition2: (m + n + 1) // 2 - partition1

   • +1 нужна для нечётных длин (чтобы больше элементов слева)

2. Граничные условия: используем float('-inf') и float('inf')

   • Если partition = 0, нет левого элемента
   • Если partition = m, нет правого элемента

3. Условие корректности: left1 <= right2 AND left2 <= right1

   • Гарантирует, что объединённый массив отсортирован

4. Расчёт медианы:

   • Чётная: берём максимум слева и минимум справа
   • Нечётная: берём максимум слева

Визуализация

nums1 = [1, 3]
nums2 = [2]

Медиана трёх чисел [1, 2, 3] — это 2

Разделение:
[1] | [3]
[2] |

sleft1=1, right1=3
left2=2, right2=inf

Условие: 1 <= inf AND 2 <= 3 ✓
Медиана (нечётная) = max(1, 2) = 2

Заключение

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