Median of Two Sorted Arrays A Deep Dive for Devs

The “Median of Two Sorted Arrays” problem is a notorious one in the world of coding interviews. It’s often presented as a test of your understanding of algorithms, specifically binary search, and your ability to handle edge cases gracefully. While a brute-force approach is straightforward, the real challenge lies in achieving an optimal O(log(min(m,n))) time complexity.

Let’s dissect this problem, starting with the basics and building up to the elegant solution.

The Problem Statement

Given two sorted arrays, nums1 of size m and nums2 of size n, return the median of the two sorted arrays.

The overall run time complexity should be O(log(m+n)) or more specifically O(log(min(m,n))).

What is a Median?

  • If the total number of elements N is odd, the median is the middle element. For example, in [1, 2, 3, 4, 5], the median is 3.
  • If the total number of elements N is even, the median is the average of the two middle elements. For example, in [1, 2, 3, 4, 5, 6], the median is (3 + 4) / 2 = 3.5.

Approach 1: Merge and Find (The Naive Way)

The most intuitive approach is to combine both sorted arrays into a single, larger sorted array, and then find the median from that merged array. This leverages the fact that both input arrays are already sorted.

Concept

  1. Create a new array, merged_array, to store elements from nums1 and nums2.
  2. Use two pointers, one for nums1 (p1) and one for nums2 (p2), to iterate through both arrays simultaneously.
  3. Compare elements at p1 and p2. Append the smaller element to merged_array and advance its respective pointer.
  4. Once one array is exhausted, append the remaining elements from the other array.
  5. Calculate the median from merged_array based on its total length.

Code Example (Python)

import math

def find_median_naive(nums1, nums2):
    """
    Finds the median of two sorted arrays by merging them.
    Time Complexity: O(m + n)
    Space Complexity: O(m + n)
    """
    m, n = len(nums1), len(nums2)
    merged_array = []
    p1, p2 = 0, 0

    # Merge the two sorted arrays
    while p1 < m and p2 < n:
        if nums1[p1] <= nums2[p2]:
            merged_array.append(nums1[p1])
            p1 += 1
        else:
            merged_array.append(nums2[p2])
            p2 += 1

    # Append remaining elements
    while p1 < m:
        merged_array.append(nums1[p1])
        p1 += 1
    while p2 < n:
        merged_array.append(nums2[p2])
        p2 += 1

    total_length = len(merged_array)

    # Calculate median
    if total_length % 2 == 1:
        # Odd number of elements
        return float(merged_array[total_length // 2])
    else:
        # Even number of elements
        mid1 = merged_array[total_length // 2 - 1]
        mid2 = merged_array[total_length // 2]
        return float(mid1 + mid2) / 2.0

print("--- Naive Approach Examples ---")
print(f"nums1 = [1, 3], nums2 = [2] -> Median: {find_median_naive([1, 3], [2])}")
print(f"nums1 = [1, 2], nums2 = [3, 4] -> Median: {find_median_naive([1, 2], [3, 4])}")
print(f"nums1 = [0, 0], nums2 = [0, 0] -> Median: {find_median_naive([0, 0], [0, 0])}")
print(f"nums1 = [], nums2 = [1] -> Median: {find_median_naive([], [1])}")
print(f"nums1 = [1], nums2 = [] -> Median: {find_median_naive([1], [])}")
print(f"nums1 = [2], nums2 = [] -> Median: {find_median_naive([2], [])}")
print(f"nums1 = [1, 2], nums2 = [3] -> Median: {find_median_naive([1, 2], [3])}")
print(f"nums1 = [1], nums2 = [2, 3] -> Median: {find_median_naive([1], [2, 3])}")

Sample Output

--- Naive Approach Examples ---
nums1 = [1, 3], nums2 = [2] -> Median: 2.0
nums1 = [1, 2], nums2 = [3, 4] -> Median: 2.5
nums1 = [0, 0], nums2 = [0, 0] -> Median: 0.0
nums1 = [], nums2 = [1] -> Median: 1.0
nums1 = [1], nums2 = [] -> Median: 1.0
nums1 = [2], nums2 = [] -> Median: 2.0
nums1 = [1, 2], nums2 = [3] -> Median: 2.0
nums1 = [1], nums2 = [2, 3] -> Median: 2.0

Analysis of Naive Approach

  • Time Complexity: O(m + n) because we iterate through both arrays once to merge them.
  • Space Complexity: O(m + n) to store the merged_array.

While correct, this approach doesn’t meet the O(log(min(m,n))) requirement. This is perfectly acceptable if you’re explicitly asked for a O(m+n) solution or if it’s a warm-up, but for the optimal solution, we need something smarter.

Approach 2: Binary Search (The Optimal Way)

This is where the problem truly shines. The O(log(min(m,n))) complexity hints at a binary search. But what are we binary searching on? Not an index in the merged array, but rather a “partition point.”

Core Idea: Finding the Perfect Partition

Imagine we split nums1 into two parts: a left part LeftX and a right part RightX. Similarly, we split nums2 into LeftY and RightY.

If we combine LeftX and LeftY to form the “left half” of the overall merged array, and RightX and RightY to form the “right half”, then for the median to be correct, two conditions must hold:

  1. Correct Number of Elements: The total number of elements in the left half (i.e., len(LeftX) + len(LeftY)) must be (m + n + 1) // 2. This ensures that we have precisely the elements forming the first half of the sorted combined array.
  2. Order Condition: max(LeftX, LeftY) must be less than or equal to min(RightX, RightY). This means the largest element in the left half is smaller than or equal to the smallest element in the right half, maintaining the sorted property.

If these two conditions are met, the median is:

  • Odd Total Length: max(max(LeftX), max(LeftY))
  • Even Total Length: (max(max(LeftX), max(LeftY)) + min(min(RightX), min(RightY))) / 2.0

Binary Search Strategy

We will binary search on the smaller of the two arrays. Let’s assume nums1 is the smaller array (if not, we’ll swap them).

  1. Choose the smaller array: Let nums1 be X and nums2 be Y. Ensure len(X) <= len(Y). If len(nums1) > len(nums2), swap them.
  2. Define Binary Search Range: We’ll binary search for the correct partition point partitionX in X. The range for partitionX will be [0, len(X)] (inclusive of 0 and len(X) to handle empty partitions).
  3. Calculate partitionY: Once partitionX is chosen, partitionY for Y is determined by the total number of elements needed in the left half. total_elements_in_left_half = (len(X) + len(Y) + 1) // 2 partitionY = total_elements_in_left_half - partitionX
  4. Identify Elements at Partitions:
    • maxLeftX: Element just before partitionX in X. (Handle partitionX == 0 with float('-inf'))
    • minRightX: Element at partitionX in X. (Handle partitionX == len(X) with float('inf'))
    • maxLeftY: Element just before partitionY in Y. (Handle partitionY == 0 with float('-inf'))
    • minRightY: Element at partitionY in Y. (Handle partitionY == len(Y) with float('inf'))
  5. Check Conditions and Adjust Search Range:
    • If maxLeftX <= minRightY AND maxLeftY <= minRightX: This is the perfect partition. We found our median.
      • If (len(X) + len(Y)) is odd, median is max(maxLeftX, maxLeftY).
      • If (len(X) + len(Y)) is even, median is (max(maxLeftX, maxLeftY) + min(minRightX, minRightY)) / 2.0.
    • If maxLeftX > minRightY: This means maxLeftX is too large, it belongs in the right half. We need to move partitionX to the left. So, high = partitionX - 1.
    • If maxLeftY > minRightX: This means maxLeftY is too large, it belongs in the right half. We need to move partitionX to the right (which will move partitionY to the left). So, low = partitionX + 1.

Note: The +1 in (len(X) + len(Y) + 1) // 2 is crucial for correctly handling both odd and even total lengths. It ensures that total_elements_in_left_half always refers to the number of elements in the first “half” of the combined array, effectively giving us the index of the first middle element for even lengths, or the single middle element for odd lengths.

Code Example (Python)

def find_median_optimal(nums1, nums2):
    """
    Finds the median of two sorted arrays using binary search on partitions.
    Time Complexity: O(log(min(m, n)))
    Space Complexity: O(1)
    """
    # Ensure nums1 is the smaller array to binary search on for efficiency
    if len(nums1) > len(nums2):
        nums1, nums2 = nums2, nums1

    m, n = len(nums1), len(nums2)
    low, high = 0, m  # Binary search range for partitionX in nums1

    # total_half_len is the number of elements required in the left partition
    # (m + n + 1) // 2 handles both odd and even total lengths correctly
    # For odd: gives (total + 1) // 2, which is the index of the median
    # For even: gives (total) // 2, which is the index of the first of two medians
    total_half_len = (m + n + 1) // 2

    while low <= high:
        partitionX = (low + high) // 2
        partitionY = total_half_len - partitionX

        # Determine elements at the boundaries of partitions
        # If partition is at the beginning, maxLeft is -infinity
        # If partition is at the end, minRight is +infinity
        maxLeftX = float('-inf') if partitionX == 0 else nums1[partitionX - 1]
        minRightX = float('inf') if partitionX == m else nums1[partitionX]

        maxLeftY = float('-inf') if partitionY == 0 else nums2[partitionY - 1]
        minRightY = float('inf') if partitionY == n else nums2[partitionY]

        # Check if we found the perfect partition
        if maxLeftX <= minRightY and maxLeftY <= minRightX:
            # We found the perfect split
            if (m + n) % 2 == 1:
                # Odd total length: median is the maximum of the left half
                return float(max(maxLeftX, maxLeftY))
            else:
                # Even total length: median is average of max of left half and min of right half
                return float(max(maxLeftX, maxLeftY) + min(minRightX, minRightY)) / 2.0
        elif maxLeftX > minRightY:
            # maxLeftX is too large, meaning partitionX is too far right.
            # We need to move partitionX to the left.
            high = partitionX - 1
        else: # maxLeftY > minRightX
            # maxLeftY is too large, meaning partitionY is too far right.
            # This implies partitionX is too far left.
            # We need to move partitionX to the right.
            low = partitionX + 1

    # Should not be reached if inputs are valid arrays as per problem constraints
    raise ValueError("Input arrays are not sorted or contain invalid data.")

print("\n--- Optimal Approach Examples ---")
print(f"nums1 = [1, 3], nums2 = [2] -> Median: {find_median_optimal([1, 3], [2])}")
print(f"nums1 = [1, 2], nums2 = [3, 4] -> Median: {find_median_optimal([1, 2], [3, 4])}")
print(f"nums1 = [0, 0], nums2 = [0, 0] -> Median: {find_median_optimal([0, 0], [0, 0])}")
print(f"nums1 = [], nums2 = [1] -> Median: {find_median_optimal([], [1])}")
print(f"nums1 = [1], nums2 = [] -> Median: {find_median_optimal([1], [])}")
print(f"nums1 = [2], nums2 = [] -> Median: {find_median_optimal([2], [])}")
print(f"nums1 = [1, 2], nums2 = [3] -> Median: {find_median_optimal([1, 2], [3])}")
print(f"nums1 = [1], nums2 = [2, 3] -> Median: {find_median_optimal([1], [2, 3])}")
print(f"nums1 = [1, 2, 3, 4, 5], nums2 = [6, 7, 8, 9, 10] -> Median: {find_median_optimal([1, 2, 3, 4, 5], [6, 7, 8, 9, 10])}")
print(f"nums1 = [1, 2, 3, 4, 5, 6], nums2 = [] -> Median: {find_median_optimal([1, 2, 3, 4, 5, 6], [])}")

Sample Output

--- Optimal Approach Examples ---
nums1 = [1, 3], nums2 = [2] -> Median: 2.0
nums1 = [1, 2], nums2 = [3, 4] -> Median: 2.5
nums1 = [0, 0], nums2 = [0, 0] -> Median: 0.0
nums1 = [], nums2 = [1] -> Median: 1.0
nums1 = [1], nums2 = [] -> Median: 1.0
nums1 = [2], nums2 = [] -> Median: 2.0
nums1 = [1, 2], nums2 = [3] -> Median: 2.0
nums1 = [1], nums2 = [2, 3] -> Median: 2.0
nums1 = [1, 2, 3, 4, 5], nums2 = [6, 7, 8, 9, 10] -> Median: 5.5
nums1 = [1, 2, 3, 4, 5, 6], nums2 = [] -> Median: 3.5

Analysis of Optimal Approach

  • Time Complexity: O(log(min(m, n))). The binary search halves the search space (which is min(m,n)) in each step. Operations inside the loop are constant time.
  • Space Complexity: O(1) as we only use a few variables to store pointers and values, no additional data structures dependent on input size.

This solution is considerably more complex to grasp initially due to the abstract nature of binary searching on partitions rather than concrete indices. However, it’s the standard solution for this problem due to its efficiency.

Key Takeaways and Reflections

  1. Complexity Matters: The naive O(m+n) solution is easy to implement but falls short for large inputs. Understanding when and why to optimize is crucial.
  2. Binary Search Power: This problem is a brilliant example of how binary search isn’t just for finding elements but also for finding “optimal states” or “split points” in a conceptually sorted space.
  3. Edge Case Handling: Using float('-inf') and float('inf') for boundary conditions (when a partition is at the very beginning or end of an array) simplifies the logic immensely, avoiding numerous if/else checks for empty left/right partitions.
  4. Conceptual Partitioning: The core idea revolves around dividing the combined elements into two halves, ensuring the elements in the left half are always less than or equal to elements in the right half.

This problem is a common litmus test for algorithm skills in interviews. While the O(log(min(m,n))) solution is tricky, understanding its mechanics provides deep insights into binary search and partitioning problems. Practice helps solidify these concepts.

Last updated on