34. Median of Two Sorted Arrays
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 == mnums2.length == n0 <= m <= 10000 <= n <= 10001 <= 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:
- Left partition has (m+n+1)//2 elements total
- 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
-
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.
-
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.
-
Partition Size: For combined length (m+n), the left partition should have
(m+n+1)//2elements. The+1handles both odd and even cases correctly. -
Edge Cases with Infinity: Using
-infand+inffor out-of-bounds partitions simplifies comparison logic without special cases. - 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)
- 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)