4. Median of Two Sorted Arrays

March 2019 4 minute read

Problem

There are two sorted arrays nums1 and nums2 of size m and n respectively.

Find the median of the two sorted arrays. The overall run time complexity should be O(log (m+n)).

You may assume nums1 and nums2 cannot be both empty.

nums1 = [1, 3]
nums2 = [2]

The median is 2.0
nums1 = [1, 2]
nums2 = [3, 4]

The median is (2 + 3)/2 = 2.5

Solution

The naive solution for this problem is merging two arrays ,and find the median of new array.

The time complexity of naive solution is O((n+m)log(m+n)).

class Solution(object):
    def findMedianSortedArrays(self, nums1, nums2):
        new = sorted(nums1+nums2)
        if len(new)%2==0:
            return (new[len(new)//2-1]+new[len(new)//2])/2
        return new[len(new)//2]

We can see that the key point to solve this problem is finding the k th element two sorted arrays fast.

Here is the algorithm for finding k th element in two sorted arrays.

def find_kth(nums1,nums2,k):
  if nums1[k/2] == nums2[k/2]:
     return nums1[k/2]
  else if nums1[k/2] < nums2[k/2]:
    return find_kth(nums1[k/2:],nums2,k/2)
  else if nums1[k/2] > nums2[k/2]:
    return find_kth(nums1,nums2[:k/2],k/2)
  1. If nums1[k/2] == nums2[k/2] ,it means there are k/2 elements greater than nums1[k/2] and k/2 elements greater than nums2[k/2]. k/2 + k/2 = k, so just return nums1[k/2]

  2. If nums1[k/2] < nums2[k/2], it means there is k/2+1 elements in nums great than nums1[k/2]. And there is k/2 elements in nums1 greater than nums1[k/2], so k th element is not in nums1[0:k/2], we can drop it.

  3. If nums1[k/2] > nums[k/2], it means there is k/2+1 elements in nums1 greater than nums2[k/2], and there is k/2 elements in nums2 greater than nums2[k/2], so k th element is not in nums2[0:k/2] , we can drop it.

Code

class Solution:
    def find_k_th(self, nums1, nums2, k):
        m = len(nums1)
        n = len(nums2)
        # expect m is equal or smaller than n
        if m > n:
            return self.find_k_th(nums2, nums1, k)
        if m == 0:
            return nums2[k - 1]
        if k == 1:
            return min(nums1[0], nums2[0])
        p1 = min(int(k / 2), m)
        p2 = k - p1
        if nums1[p1 - 1] <= nums2[p2 - 1]:
            return self.find_k_th(nums1[p1:], nums2, p2)
        else:
            return self.find_k_th(nums1, nums2[p2:], p1)

    def findMedianSortedArrays(self, nums1: List[int],
                               nums2: List[int]) -> float:
       length = len(nums1) + len(nums2)
        mid = length >> 1
        mid_1 = self.find_k_th(nums1, nums2, mid + 1)
        if length % 2 != 0:
            return mid_1
        return (mid_1 + self.find_k_th(nums1, nums2, mid)) / 2
int findKth(vector<int> &a, int a_start, int a_end, vector<int> &b, int b_start, int b_end, int k)
{
    // corner case
    if (a_end - a_start > b_end - b_start)
    {
        return findKth(b, b_start, b_end, a, a_start, a_end, k);
    }
    if (a_end < a_start)
    {
        return b[k - 1];
    }
    if (k == 1)
    {
        return min(a[a_start], b[b_start]);
    }

    int pa = min(k / 2, a_end - a_start + 1);
    int pb = k - pa;
    if (a[a_start + pa - 1] == b[b_start + pb - 1])
    {
        return a[a_start + pa - 1];
    }
    if (a[a_start + pa - 1] < b[b_start + pb - 1])
    {
        return findKth(a, a_start + pa, a_end, b, b_start, b_end, k - pa);
    }
    return findKth(a, a_start, a_end, b, b_start + pa, b_end, k - pa);
}

class Solution
{
  public:
    double findMedianSortedArrays(vector<int> &nums1, vector<int> &nums2)
    {
        int m = nums1.size();
        int n = nums2.size();
        int total = m + n;
        double a = (double)findKth(nums1, 0, m - 1, nums2, 0, n - 1, total / 2 + 1);
        if (total % 2)
        {
            return a;
        }
        int b = findKth(nums1, 0, m - 1, nums2, 0, n - 1, total / 2);
        return (a + b) * 0.5;
    }
};
func min(x, y int) int {
    if x > y {
        return y
    }
    return x
}

func findKth(nums1, nums2 []int, k int) int {
    m := len(nums1)
    n := len(nums2)
    // we assume that len(nums2) is greater than len(nums1)
    if m > n {
        return findKth(nums2, nums1, k)
    }
    if m == 0 {
        return nums2[k-1]
    }
    if k == 1 {
        return min(nums2[0], nums1[0])
    }
    p1 := min(k>>1, m)
    p2 := k - p1
    if nums1[p1-1] == nums2[p2-1] {
        return nums1[p1-1]
    } else if nums1[p1-1] < nums2[p2-1] {
        return findKth(nums1[p1:], nums2, p2)
    }
    return findKth(nums1, nums2[p2:], p1)
}

func findMedianSortedArrays(nums1 []int, nums2 []int) float64 {
    m := len(nums1)
    n := len(nums2)
    // even number
    if (m+n)%2 == 0 {
        return float64(findKth(nums1, nums2, ((m+n)>>1)+1)+findKth(nums1, nums2, (m+n)>>1)) * 0.5
    }
    return float64(findKth(nums1, nums2, (m+n)>>1+1))
}