Median of Two Sorted Arrays

Problem Id: 4 Difficulty: Hard Tag: Array Tag: Binary Search Tag: Divide and Conquer


Intuition

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 .

Solution


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