View on NeetCode
View on LeetCode

Problem

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

The overall run time complexity should be O(log (m+n)).

Example 1:

Input: nums1 = [1,3], nums2 = [2]
Output: 2.00000
Explanation: merged array = [1,2,3] and median is 2.

Example 2:

Input: nums1 = [1,2], nums2 = [3,4]
Output: 2.50000
Explanation: merged array = [1,2,3,4] and median is (2 + 3) / 2 = 2.5.

Constraints:

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

Solution

Approach: Binary Search on Smaller Array

The key insight is to partition both arrays such that:

  1. Left partition has (m+n+1)//2 elements total
  2. All elements in left partition <= all elements in right partition

We binary search on the smaller array to find the correct partition point. For a valid partition:

  • max(left_nums1, left_nums2) <= min(right_nums1, right_nums2)

The median is:

  • If total length is odd: max(left_nums1, left_nums2)
  • If even: (max(left_nums1, left_nums2) + min(right_nums1, right_nums2)) / 2

Implementation

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        # Ensure nums1 is the smaller array
        if len(nums1) > len(nums2):
            nums1, nums2 = nums2, nums1

        m, n = len(nums1), len(nums2)
        left, right = 0, m

        while left <= right:
            partition1 = (left + right) // 2
            partition2 = (m + n + 1) // 2 - partition1

            # Handle edge cases with infinity
            maxLeft1 = float('-inf') if partition1 == 0 else nums1[partition1 - 1]
            minRight1 = float('inf') if partition1 == m else nums1[partition1]

            maxLeft2 = float('-inf') if partition2 == 0 else nums2[partition2 - 1]
            minRight2 = float('inf') if partition2 == n else nums2[partition2]

            # Check if we found the correct partition
            if maxLeft1 <= minRight2 and maxLeft2 <= minRight1:
                # Found the correct partition
                if (m + n) % 2 == 0:
                    return (max(maxLeft1, maxLeft2) + min(minRight1, minRight2)) / 2
                else:
                    return max(maxLeft1, maxLeft2)
            elif maxLeft1 > minRight2:
                # Too far right in nums1, go left
                right = partition1 - 1
            else:
                # Too far left in nums1, go right
                left = partition1 + 1

        # Should never reach here if inputs are valid
        return 0.0

Complexity Analysis

  • Time Complexity: O(log(min(m, n))), where m and n are the lengths of the two arrays. We perform binary search on the smaller array.
  • Space Complexity: O(1), we only use a constant amount of extra space.

Key Insights

  1. Partition Strategy: Instead of merging arrays, we partition both arrays such that all elements on the left are smaller than all elements on the right. This gives us the median elements at the partition boundary.

  2. Binary Search on Smaller Array: We always search on the smaller array to minimize the search space. This ensures O(log(min(m, n))) time complexity.

  3. Partition Size: For combined length (m+n), the left partition should have (m+n+1)//2 elements. The +1 handles both odd and even cases correctly.

  4. Edge Cases with Infinity: Using -inf and +inf for out-of-bounds partitions simplifies comparison logic without special cases.

  5. Validation Condition: A valid partition satisfies:
    • maxLeft1 <= minRight2 (largest in left half of nums1 <= smallest in right half of nums2)
    • maxLeft2 <= minRight1 (largest in left half of nums2 <= smallest in right half of nums1)
  6. Median Calculation:
    • Odd total length: median is the max of the left partition
    • Even total length: median is the average of max(left) and min(right)