0004. Median of Two Sorted Arrays

0004. Median of Two Sorted Arrays#

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))\).

Examples#

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

Analysis#

  • solution 1: this is not a sollution because the complexity is \(O(m+n)\)

    1. Merge two sorted arrays into an additional array

    2. find the median of the additional array

    3. how to design the solution that uses \(O(m+n)\) space?

  • solution 2:

    • median is the middle number in a sorted list of numbers. The rest assumes ascending lists.

    • The median partitions the list into two parts: the left partition is no greater than the median and the right partition is no smaller than the median.

    • So we can use binary search to find such an index that the maximum number in the left partition is no greater than the minimum number in the right partition.

Solution#

# Solution 1: naive solution O(m+n) - two-pointers
def solution(nums1, nums2):
    m = len(nums1)
    n = len(nums2)
    i = 0
    j = 0
    k = 0
    nums = [0]*(m+n)
    # use merge-sort 
    # if m!=n, one of the array will be first fully traversed
    while i < m and j < n:
        if nums1[i] < nums2[j]:
            nums[k] = nums1[i]
            i += 1
        else:
            nums[k] = nums2[j]
            j += 1
        k += 1
    # after one of the array is traversed, need traverse the other array for the remaining elements
    while i < m:
        nums[k] = nums1[i]
        i += 1
        k += 1
    
    while j < n:
       nums[k] = nums2[j]
       j += 1
       k += 1 
    
    # find the median
    median_index = (m+n) // 2
    if (m+n) % 2 == 1:
        median = nums[median_index]
    else:
        median = (nums[median_index] + nums[median_index-1])/2.

    return median_index, median

# test
nums1 = [1, 3]
nums2 = [2]
print(solution(nums1, nums2))

nums1 = [1, 2]
nums2 = [3, 4]
print(solution(nums1, nums2))

(1, 2)
(2, 2.5)

def solution(nums1, nums2):
    A, B = nums1, nums2
    total = len(nums1) + len(nums2)
    half = total // 2
    
    # start with the shortest array
    if len(B) < len(A):
        A, B = B, A 
    
    # left and right pointer for A
    l, r = 0, len(A) - 1
    while True:
        i = (l + r) // 2 # A: half index for A to make the first part of the left partition
        j = half - i - 2 # B: the rest of numbers needed to make up the left partition from 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")

        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

# test
nums1 = [1, 3]
nums2 = [2]
print(solution(nums1, nums2))

nums1 = [1, 2]
nums2 = [3, 4]
print(solution(nums1, nums2))

2
2.5
Contents