To get the median, we need to split the 2 arrays both into 2 parts. And these 2 parts should follow conditions below:
len(left_part) = len(right_part) or len(right_part) + 1 max(left_part) <= min(right_part)
With these 2 parts, we could simply get the median, based on the total number are odd or even.
Let's say nums1 has m numbers, num2 has n numbers and m <= n . So our purpose is to find i from [0, m] to split nums1 into nums1[:i] and nums1[i:]. Since 2 parts have the same number of values.:
i + j = m + n - i - j [+ 1]
So,:
j = (m + n + 1) // 2 - i
Another condition we need to consider is max(left_part) <= min(right_part) . And we would use this to do the binary search.
nums1[i - 1] <= nums2[j] and nums2[j - 1] <= nums1[i]
If nums1[i - 1] > nums2[j] , i is too big, so we could set hi to i - 1 .
Else if nums2[j - 1] > nums1[i], j is too big, and i is too small. So we could set lo to i + 1 .
class Solution:
def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
m, n = len(nums1), len(nums2)
if m > n:
nums1, nums2, m, n = nums2, nums1, n, m
lo = 0
hi = m
while True:
i = (lo + hi) // 2
j = (m + n + 1) // 2 - i
if i > 0 and nums1[i - 1] > nums2[j]:
hi = i - 1
elif i < m and nums2[j - 1] > nums1[i]:
lo = i + 1
else:
# Have found the correct i
if i == 0:
left_max = nums2[j - 1]
elif j == 0:
left_max = nums1[i - 1]
else:
left_max = max(nums1[i - 1], nums2[j - 1])
if (m + n) % 2 == 1:
return float(left_max)
if i == m:
right_min = nums2[j]
elif j == n:
right_min = nums1[i]
else:
right_min = min(nums1[i], nums2[j])
return (right_min + left_max) / 2