4. Median of Two Sorted Arrays — LeetCode(Python)
I got you!
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
-106 <= nums1[i], nums2[i] <= 106
Solution:
A, B = nums1, nums2
total = len(nums1) + len(nums2)
half = total // 2
if len(B) < len(A):
A, B = B, A
l, r = 0, len(A) - 1
while True:
i = (l + r) // 2 # A
j = half - i - 2 # B
Aleft = A[i] if i >= 0 else float("-infinity")
Aright = A[i + 1] if (i + 1) < len(A) else float("infinity")
Bleft = B[j] if j >= 0 else float("-infinity")
Bright = B[j + 1] if (j + 1) < len(B) else float("infinity")
# partition is correct
if Aleft <= Bright and Bleft <= Aright:
# odd
if total % 2:
return min(Aright, Bright)
# even
return (max(Aleft, Bleft) + min(Aright, Bright)) / 2
elif Aleft > Bright:
r = i - 1
else:
l = i + 1
Explanation —
We perform binary search on the smaller array here, which we have set to A, and the other to B.
At every iteration of our binary search procedure —
- we look for the mid element of A at index i
- we then store the mid index of B at j, which is just (half - i)[-2 for indexing correctly]
- Aleft is the rightmost element of the left subarray of A, Aright is the leftmost element of the right subarray of A. Similarly, we have Bleft and Bright.
- If the partition between A and B is correct, we calculate the median and return it.
- If Aleft > Bright, it means we are ahead the optimal index in A, so we shift the right end to left of mid to calculate a new mid(i) in the next iteration.
- Else if Bleft > Aright, it means we are behind the optimal index in A, so we shift our left pointer to find a new mid(i).
Time and Space Complexity —
Since we are performing binary search on array A and doing a constant time operation to calculate our mid for B, we can say that our algorithm requires logarithmic time to run.
Also, we require a constant amount of extra space.
Thus, if m and n are the lengths of the two arrays we are given,
Time Complexity: O(min(m, n))
Space Complexity: O(1)