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 is3
. - 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
- Create a new array,
merged_array
, to store elements fromnums1
andnums2
. - Use two pointers, one for
nums1
(p1
) and one fornums2
(p2
), to iterate through both arrays simultaneously. - Compare elements at
p1
andp2
. Append the smaller element tomerged_array
and advance its respective pointer. - Once one array is exhausted, append the remaining elements from the other array.
- 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 themerged_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:
- 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. - Order Condition:
max(LeftX, LeftY)
must be less than or equal tomin(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).
- Choose the smaller array: Let
nums1
beX
andnums2
beY
. Ensurelen(X) <= len(Y)
. Iflen(nums1) > len(nums2)
, swap them. - Define Binary Search Range: We’ll binary search for the correct partition point
partitionX
inX
. The range forpartitionX
will be[0, len(X)]
(inclusive of0
andlen(X)
to handle empty partitions). - Calculate
partitionY
: OncepartitionX
is chosen,partitionY
forY
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
- Identify Elements at Partitions:
maxLeftX
: Element just beforepartitionX
inX
. (HandlepartitionX == 0
withfloat('-inf')
)minRightX
: Element atpartitionX
inX
. (HandlepartitionX == len(X)
withfloat('inf')
)maxLeftY
: Element just beforepartitionY
inY
. (HandlepartitionY == 0
withfloat('-inf')
)minRightY
: Element atpartitionY
inY
. (HandlepartitionY == len(Y)
withfloat('inf')
)
- Check Conditions and Adjust Search Range:
- If
maxLeftX <= minRightY
ANDmaxLeftY <= minRightX
: This is the perfect partition. We found our median.- If
(len(X) + len(Y))
is odd, median ismax(maxLeftX, maxLeftY)
. - If
(len(X) + len(Y))
is even, median is(max(maxLeftX, maxLeftY) + min(minRightX, minRightY)) / 2.0
.
- If
- If
maxLeftX > minRightY
: This meansmaxLeftX
is too large, it belongs in the right half. We need to movepartitionX
to the left. So,high = partitionX - 1
. - If
maxLeftY > minRightX
: This meansmaxLeftY
is too large, it belongs in the right half. We need to movepartitionX
to the right (which will movepartitionY
to the left). So,low = partitionX + 1
.
- If
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 ismin(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
- 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. - 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.
- Edge Case Handling: Using
float('-inf')
andfloat('inf')
for boundary conditions (when a partition is at the very beginning or end of an array) simplifies the logic immensely, avoiding numerousif/else
checks for empty left/right partitions. - 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.