Problem Description

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.

My Idea

The idea here is to separate both arrays A and B in two partitions left and right, so that the length of A's right partition + B's right partition = 1/2 * len(A+B) and all elements in A's left partition are <= than all elements in B's right partition AND all elements in B's left partition are <= than all elements in A's right partition. This means that A's left partition + B's left partition is essentially the left partition of the combined array and the same goes for the right partitions respectively. To do this we run a sort of binary search on the smaller array A to minimize the search space and partition B respectively so we fulfill the length requirements, while adjusting A's binary search pointers in accordance to the requirement for the values of the partitions. The edge cases are covered using infinite bounds to spare more complex checks. In the end this approach yields a time complexity of O(log min(n,m)), since we’re essentially running a binary search on the smaller array.

My solution

# Time complexity: O(log min(n,m))
def findMedianSortedArrays(nums1: list[int], nums2: list[int]) -> float:
    # Name the arrays and make A be the shorter one
    if len(nums1)<len(nums2):
        A = nums1
        B = nums2
    else:
        A = nums2
        B = nums1
    # Total length of merged array
    total = len(A) + len(B)
    # Half of the length of merged array
    half = total // 2
    # initiate pointers for binary search in A
    l = 0
    r = len(A)-1

    while True:
        i = (l + r) // 2 # mid for A
        j = half - i - 2 # boundry for B, so that len(A[0:i+1]) + len(B[0:j+1]) = half

        # Set pointers for the boundries of A and B (handling out of bound indeces with '+/-inf')
        Aleft = A[i] if i >= 0 else float('-inf')
        Aright = A[i+1] if (i+1)<len(A) else float('inf')
        Bleft = B[j] if j >= 0 else float('-inf')
        Bright = B[j+1] if (j+1)<len(B) else float('inf')

        # Partition is correct
        if Aleft <= Bright and Bleft <=Aright:
            # odd number of elements -> return the min of the right partitions
            if total % 2 :
               return min(Aright,Bright)
            # even number -> return the mean of the max of left and min of right partition
            return (max(Aleft,Bleft) + min (Aright,Bright)) / 2
        # A's left partition is too big
        elif Aleft > Bright:
            r = i - 1
        # A's left partition is too small
        else:
            l = i + 1