leetcode - 4. Median of Two Sorted Arrays

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.

Constraints:

nums1.length == m
nums2.length == n
0 <= m <= 1000
0 <= n <= 1000
1 <= m + n <= 2000
-10^6 <= nums1[i], nums2[i] <= 10^6

Solution

Merge Sort

Calculate the final position of the median, and merge 2 sorted list to find out the median.

Time complexity: $o(n+m)$
Space complexity: $o(1)$

Binary Search

Refer to this video.

We have 2 lists X, Y as below:

		Px
		|
X = [x1, x2, x3, x4]
Y = [y1, y2, y3, y4, y5, y6]
					|
					Py

And we want to find a partition Px and a partition Py, where there are equal numbers at the left and right of the partition. As showed above, there are 1 element at the left of Px, and 4 elements at the left of Py, which means 5 elements at the left of partition and 5 elements at the right.

If we can find the partition, and x1 <= y5, y4 <= x2, that means all the elements at the left of the partition are smaller than the right, then we found our median.

Use binary search to find partition x, and then partition y will be: (n1 + n2 + 1)//2.

Some corner cases:

• If the number of merged lists is odd, then we have 1 more element at the left.

• If the number of merged lists is odd, then the median is max(x1, y4)

• If the number of merged lists is even, the the median is avg(max(x1, y4), min(x2, y5))

• For partition, we might have empty left or right. For example if we have the partition like this:

    	  Px
    	  |
    X = [x1, x2, x3, x4]
    Y = [y1, y2, y3, y4, y5, y6]
    						  |
    						  Py
  

  Then, we use float(-inf) to denote the number at the left of Px, similarly, float(inf) to denote the number of right half of partition.

• The exit of binary becomes left <= right.

Code

Merge Sort

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m, n = len(nums1), len(nums2)
        final_len = m + n
        p1, p2 = 0, 0
        p = 0
        cur_num = None
        candidates = []
        while p1 < m or p2 < n:
            if p == final_len // 2 + 1:
                candidates.append(cur_num)
            if p == final_len // 2 and final_len % 2 == 0:
                candidates.append(cur_num)
            if p1 < m and p2 < n:
                if nums1[p1] <= nums2[p2]:
                    cur_num = nums1[p1]
                    p1 += 1
                else:
                    cur_num = nums2[p2]
                    p2 += 1
            elif p1 < m:
                cur_num = nums1[p1]
                p1 += 1
            elif p2 < n:
                cur_num = nums2[p2]
                p2 += 1
            p += 1
        if p == final_len // 2 + 1:
            candidates.append(cur_num)
        if p == final_len // 2 and final_len % 2 == 0:
            candidates.append(cur_num)            
        return sum(candidates) / len(candidates)

Binary Search

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        if len(nums1) < len(nums2):
            x, y = nums1, nums2
        else:
            x, y = nums2, nums1
        n = len(nums1) + len(nums2)
        half_n = (n + 1) // 2
        left, right = 0, len(x)
        while left <= right:
            px = (left + right) >> 1
            py = half_n - px
            x_left = x[px - 1] if px > 0 else float('-inf')
            x_right = x[px] if px < len(x) else float('inf')
            y_left = y[py - 1] if py > 0 else float('-inf')
            y_right = y[py] if py < len(y) else float('inf')
            if x_left <= y_right and y_left <= x_right:
                if n & 1 == 1:
                    return float(max(x_left, y_left))
                else:
                    return (max(x_left, y_left) + min(x_right, y_right)) / 2
            elif x_left > y_right:
                right = px - 1
            elif y_left > x_right:
                left = px + 1