Median of Two Sorted Arrays

Problem Statement

You are given two integer arrays nums1 and nums2 of size m and n respectively, where each is sorted in ascending order. Return the median value among all elements of the two arrays.

Your solution must run in O(log (m+n)) time.

Example 1:

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


Output: 2.0

Explanation: Among [1, 2, 3] the median is 2.

Example 2:

Input: nums1 = [1,3], nums2 = [2,4]


Output: 2.5

Explanation: Among [1, 2, 3, 4] the median is (2 + 3) / 2 = 2.5.

Constraints:

  • nums1.length == m
  • nums2.length == n
  • 0 <= m <= 1000
  • 0 <= n <= 1000
  • -10^6 <= nums1[i], nums2[i] <= 10^6

You should aim for a solution with O(log(min(n, m))) time and O(1) space, where n is the size of nums1 and m is the size of nums2.

You should aim for a solution with O(log(min(n, m))) time and O(1) space, where n is the size of nums1 and m is the size of nums2.

Hint 1

A brute force solution would be to create a new array by merging elements from both arrays, then sorting it and returning the median. This would be an O(n + m) solution. Can you think of a better way? Maybe you can use the criteria of both the arrays being sorted in ascending order.

Hint 2

Suppose we merged both arrays. Then, we would have half = (m + n) / 2 elements to the left of the median. So, without merging, is there any way to use this information to find the median? You can leverage the fact that the arrays are sorted. Consider the smaller array between the two and use binary search to find the correct partition between the two arrays, which will allow you to directly find the median without fully merging the arrays. How will you implement this?

Hint 3

We will always try to keep array A smaller and interchange it with array B if len(A) > len(B). Now, we perform binary search on the number of elements we will choose from array A. It is straightforward that when we choose x elements from array A, we have to choose half - x elements from array B. But we should also ensure that this partition is valid. How can we do this?

Hint 4

When we do a partition for both arrays, we should ensure that the maximum elements from the left partitions of both arrays are smaller than or equal to the minimum elements of the right partitions of both the arrays. This will ensure that the partition is valid, and we can then find the median. We can find the min or max of these partitions in O(1) as these partitions are sorted in ascending order. Why does this work?

Hint 5

For example, consider the arrays A = [1, 2, 3, 4, 5] and B = [1, 2, 3, 4, 5, 6, 7, 8]. When we select x = 2, we take 4 elements from array B. However, this partition is not valid because value 4 from the left partition of array B is greater than the value 3 from the right partition of array A. So, we should try to take more elements from array A to make the partition valid. Binary search will eventually help us find a valid partition.

Solution

Brute Force

def find_median_sorted_arrays(nums1: list[int], nums2: list[int]) -> float:
    merged = sorted(nums1 + nums2)
    total_len = len(merged)
    if total_len % 2 == 0:
        return (merged[total_len // 2 - 1] + merged[total_len // 2]) / 2.0
    else:
        return float(merged[total_len // 2])

Time complexity: O((n + m)\log (n + m)) Space complexity: O(n + m)

Two Pointers

def find_median_sorted_arrays(nums1: list[int], nums2: list[int]) -> float:
    len1, len2 = len(nums1), len(nums2)
    i = j = 0
    median1 = median2 = 0


    for count in range((len1 + len2) // 2 + 1):
        median2 = median1
        if i < len1 and j < len2:
            if nums1[i] > nums2[j]:
                median1 = nums2[j]
                j += 1
            else:
                median1 = nums1[i]
                i += 1
        elif i < len1:
            median1 = nums1[i]
            i += 1
        else:
            median1 = nums2[j]
            j += 1


    if (len1 + len2) % 2 == 1:
        return float(median1)
    else:
        return (median1 + median2) / 2.0

Time complexity: O(n + m) Space complexity: O(1)

def get_kth(a: list[int], m: int, b: list[int], n: int, k: int, a_start: int = 0, b_start: int = 0) -> int:
    if m > n:
        return get_kth(b, n, a, m, k, b_start, a_start)
    if m == 0:
        return b[b_start + k - 1]
    if k == 1:
        return min(a[a_start], b[b_start])


    i = min(m, k // 2)
    j = min(n, k // 2)


    if a[a_start + i - 1] > b[b_start + j - 1]:
        return get_kth(a, m, b, n - j, k - j, a_start, b_start + j)
    else:
        return get_kth(a, m - i, b, n, k - i, a_start + i, b_start)


def find_median_sorted_arrays(nums1: list[int], nums2: list[int]) -> float:
    left = (len(nums1) + len(nums2) + 1) // 2
    right = (len(nums1) + len(nums2) + 2) // 2
    return (get_kth(nums1, len(nums1), nums2, len(nums2), left) +
            get_kth(nums1, len(nums1), nums2, len(nums2), right)) / 2.0

Time complexity: O(\log (m + n)) Space complexity: O(\log (m + n))

Binary Search (Optimal)

def find_median_sorted_arrays(nums1: list[int], nums2: list[int]) -> float:
    a, b = nums1, nums2
    total = len(nums1) + len(nums2)
    half = total // 2


    if len(b) < len(a):
        a, b = b, a


    l, r = 0, len(a) - 1
    while True:
        i = (l + r) // 2
        j = half - i - 2


        a_left = a[i] if i >= 0 else float("-infinity")
        a_right = a[i + 1] if (i + 1) < len(a) else float("infinity")
        b_left = b[j] if j >= 0 else float("-infinity")
        b_right = b[j + 1] if (j + 1) < len(b) else float("infinity")


        if a_left <= b_right and b_left <= a_right:
            if total % 2:
                return min(a_right, b_right)
            return (max(a_left, b_left) + min(a_right, b_right)) / 2
        elif a_left > b_right:
            r = i - 1
        else:
            l = i + 1

Time complexity: O(\log (min(n, m))) Space complexity: O(1)