Median of Two Sorted Arrays

Solution

Let $m$ be the length of nums1 and $n$ be the length of nums2.

Brute Force

Merge nums1 and nums2 into sorted nums, and find the median.

Complexity

time: $O(m + n)$
space: $O(m + n)$

def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        # merge 2 sorted arr
        # return the median

        p1 = p2 = 0
        n1 = len(nums1)
        n2 = len(nums2)
        mid = (n1 + n2) // 2
        arr = []

        while p1 != n1 or p2 != n2:
            if p2 == n2 or (p1 < n1 and nums1[p1] < nums2[p2]):
                arr.append(nums1[p1])
                p1 += 1
            else:
                arr.append(nums2[p2])
                p2 += 1
        
        if (n1 + n2) % 2 == 0:
            return (arr[mid] + arr[mid-1]) / 2
        else:
            return arr[mid]

Follow Up

The overall run time complexity should be $O(\log (m+n))$.

Binary Search

This video (opens new window) and This post in Chinese (opens new window) explains well.

Complexity

time: $O(\log(\min(m, n)))$
space: $O(1)$

def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
    m = len(nums1)
    n = len(nums2)
    if m > n:
        m, n = n, m
        A, B = nums2, nums1
    else:
        A, B = nums1, nums2

    imin = 0
    imax = m
    half_len = (m+n+1) // 2
    while imin <= imax:
        i = (imin + imax) // 2
        j = half_len - i
        if i < m and B[j-1] > A[i]:
            imin = i + 1
        elif i > 0 and A[i-1] > B[j]:
            imax = i - 1
        else:
            # finished searching
            if i == 0:
                max_left = B[j-1]
            elif j == 0:
                max_left = A[i-1]
            else:
                max_left = max(A[i-1], B[j-1])

            if (m+n) % 2 == 1:
                return max_left

            if i == m:
                min_right = B[j]
            elif j == n:
                min_right = A[i]
            else:
                min_right = min(A[i], B[j])
            
            return (max_left + min_right) / 2

TODO redo:

def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
    m, n = len(nums1), len(nums2)
    # nums1 should be shorter than nums2, s.t. `partitionY` is always nonnegative
    if n < m:
        nums1, nums2 = nums2, nums1
        m, n = n, m
    low, high = 0, len(nums1)
    while low <= high:
        partitionX = (low + high) // 2
        partitionY = (m + n + 1) // 2 - partitionX
        leftMaxX = -sys.maxsize if partitionX == 0 else nums1[partitionX - 1]
        rightMinX = sys.maxsize if partitionX == m else nums1[partitionX]
        leftMaxY = -sys.maxsize if partitionY == 0 else nums2[partitionY - 1]
        rightMinY = sys.maxsize if partitionY == n else nums2[partitionY]
        # partitionX found, return the median
        if leftMaxX <= rightMinY and leftMaxY <= rightMinX:
            # if the total length is odd, then return the left max
            if (m + n) % 2:
                return max(leftMaxX, leftMaxY)
            # if the total length is even, then return the average of the left max and right min
            else:
                return (max(leftMaxX, leftMaxY) + min(rightMinX, rightMinY)) / 2
        # partitionX is too much to the right
        elif leftMaxX > rightMinY:
            high = partitionX - 1
        # partitionX is too much to the left
        else:
            low = partitionX + 1
    raise Exception("Arguments were not sorted!")