Compute median of two sorted arrays
Company: Rippling
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Technical Screen
Quick Answer: This question evaluates algorithmic problem-solving with sorted arrays, focusing on concepts like binary search, partitioning strategies, complexity analysis, and careful handling of edge cases when computing order statistics.
Constraints
- 0 <= len(nums1), len(nums2) <= 10^5
- 1 <= len(nums1) + len(nums2) <= 2 * 10^5
- -10^6 <= nums1[i], nums2[i] <= 10^6
- nums1 and nums2 are each sorted in non-decreasing order
Examples
Input: ([1, 3], [2])
Expected Output: 2.0
Explanation: The combined sorted array is [1, 2, 3], so the middle element is 2.
Input: ([1, 2], [3, 4])
Expected Output: 2.5
Explanation: The combined sorted array is [1, 2, 3, 4], so the median is (2 + 3) / 2 = 2.5.
Input: ([], [7])
Expected Output: 7.0
Explanation: Only one element exists in total, so that element is the median.
Input: ([0, 0], [0, 0])
Expected Output: 0.0
Explanation: The combined sorted array is [0, 0, 0, 0], and the two middle elements are both 0.
Input: ([-5, -3, -1], [2, 4, 6, 8])
Expected Output: 2.0
Explanation: The combined sorted array is [-5, -3, -1, 2, 4, 6, 8], so the middle element is 2.
Input: ([1, 2, 2], [2, 2, 3])
Expected Output: 2.0
Explanation: The combined sorted array is [1, 2, 2, 2, 2, 3], and the two middle elements are both 2.
Solution
def solution(nums1, nums2):
# Binary search on the smaller array.
# We choose partitions i and j such that:
# - i elements are taken from nums1 into the left half
# - j elements are taken from nums2 into the left half
# - i + j == (m + n + 1) // 2
# A valid partition satisfies:
# left1 <= right2 and left2 <= right1
# Once found, the median comes from the boundary values.
if len(nums1) > len(nums2):
nums1, nums2 = nums2, nums1
m, n = len(nums1), len(nums2)
total_left = (m + n + 1) // 2
lo, hi = 0, m
while lo <= hi:
i = (lo + hi) // 2
j = total_left - i
left1 = float('-inf') if i == 0 else nums1[i - 1]
right1 = float('inf') if i == m else nums1[i]
left2 = float('-inf') if j == 0 else nums2[j - 1]
right2 = float('inf') if j == n else nums2[j]
if left1 <= right2 and left2 <= right1:
if (m + n) % 2 == 1:
return float(max(left1, left2))
return (max(left1, left2) + min(right1, right2)) / 2.0
elif left1 > right2:
hi = i - 1
else:
lo = i + 1
raise ValueError('Input arrays are not sorted or constraints are violated.')Time complexity: O(log(min(m, n))). Space complexity: O(1).
Hints
- Try binary searching on the smaller array instead of thinking about positions in the merged array directly.
- Think of splitting both arrays into a left half and a right half such that every value on the left is less than or equal to every value on the right.