LeetCode 04 Median of two Sorted Arrays

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

Example 1:

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

The median is 2.0

Example 2:

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

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

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

Difficulty
Hard

Related Problems

Analysis

This is a classic problem. It has a more general form, that is "Given two sorted arrays, find the kth element". The median of a sequence of numbers is that after the numbers are sorted, it is greater than half of the numbers, and smaller the other half.

1. O(m + n)

It's natural that we merge two sorted arrays first, then find the kth element. This will take O(m + n) time. Since we need only the kth element, sortedly merge all elements could be avoided. Use two indexers p1 and p2 pointing to the first element of array A and array B respectively. Use a counter m to record how far the pointers have gone so far. Move the indexer when the value it points to is less. When counter m equals k, we find the solution. This takes O(k) time, and O(1) space in best case. when k is close to m + n, it will take O(m+n) for worst case. Overall, it's a O(m + n) algorithm.

2. O(log(m+n))

The flaw of the above algorithm is that we rule out only one element at a time, causing time complexity up to O(m + n).
If we take advantage of binary search idea for sorted sequence, by which means rule out half number of the array at a time,
we should be able to reduce the time complexity to log(m+n).

Suppose the element number of array A and array B are both greater than k/2. We compare the k/2-th elment of A and B.
There should be 3 situations would appear:

• A[k/2-1] < B[k/2-1]
• A[k/2-1] > B[k/2-1]
• A[k/2-1] == B[k/2-1]

(1) A[k/2-1] < B[k/2-1] If we merge two arrays, A[0 ... k/2-1] will be in the former k elements. We rule them out, A becomes A[k/2-1 ... n-1]. Recursively find the target element in the A[k/2-1 ... n-1] and B.

(2) A[k/2-1] > B[k/2-1] If we merge two arrays, B[0 ... k/2-1] will be in the former k elements. We rule them out, B becomes B[k/2-1 ... n-1]. Recursively find the target element in the A and B[k/2-1 ... n-1].

(3) A[k/2-1] == B[k/2-1]
Obviously, we found the kth element, return either A[k/2-1] or B[k/2-1].

Recursion stop condition

• if A or B is empty array, return A[k-1] or B[k-1] directly.
• if A[k/2-1] == B[k/2-1], return either A[k/2-1] or B[k/2-1]
• if the recursion reduce k to 1 (k == 1), return min(A[0], B[0])

In all other cases, we call our algorithm recursively.

Note
When total number of elements is odd (e.g. k = 2n + 1), the median is nums[(2n + 1)/2 + 1], this is nums[k/2 + 1].
When total number of elements is even (e.g. k = 2n ), the median is (nums[2n/2 ] + nums[2n/2 + 1])/2, this is (nums[k/2] + nums[k/2+1])/2.

In recursion procedure, we alway put shorter array as array A, and longer array as array B.

Also, Be care when the length of the shorter array is smaller than k/2.
e.g.
if we take ia = min(k/2, m) number of elements of array A, then we should take ib = k - ia number of elements of array B.

Solutions

1. Cpp solution 35ms

class Solution {
public:
    double findMedianSortedArrays(const vector<int> &A, const vector<int> &B){
        const int m = A.size(), n = B.size();
        int total = m + n;
        if(total & 0x1) // odd
            return find_kth(A.begin(), A.end(), B.begin(), B.end(), total / 2 + 1);
        else // even
            return (find_kth(A.begin(), A.end(), B.begin(), B.end(), total / 2 )
                + find_kth(A.begin(), A.end(), B.begin(), B.end(), total / 2 + 1)) / 2.0;
    }

private:
    typedef vector<int>::const_iterator Iter;
    static int find_kth(Iter beginA, Iter endA, Iter beginB, Iter endB, int k){
        // always assume that m <= n
        const int m = distance(beginA, endA), n = distance(beginB, endB);
        if(m > n) return find_kth(beginB, endB, beginA, endA, k);
        if(m == 0) return *(beginB + k - 1);
        if(k == 1) return min(*beginA, *beginB);

        // divide k into two parts
        int ia = min(k / 2, m), ib = k - ia;
        if(*(beginA + ia - 1) < *(beginB + ib - 1)){
            return find_kth(beginA + ia, endA, beginB, endB, k - ia);
        }
        else if(*(beginA + ia - 1) > *(beginB + ib - 1)){
            return find_kth(beginA, endA, beginB + ib, endB, k - ib);
        }
        else
            return *(beginA + ia - 1);
    }
};

1. Go solution 36ms

func findMedianSortedArrays(A []int, B []int) float64 {
    m, n := len(A), len(B)
    total := m + n
    if total&0x1 == 1 { // odd
        return findKth(A, B, 0, 0, total/2+1)
    } else {
        return (findKth(A, B, 0, 0, total/2) + findKth(A, B, 0, 0, total/2+1)) / 2
    }
}

func findKth(A, B []int, beginA, beginB, k int) float64 {
    m, n := len(A)-beginA, len(B)-beginB
    if m > n {
        return findKth(B, A, beginB, beginA, k)
    }
    if m == 0 {
        return float64(B[beginB+k-1])
    }
    if k == 1 {
        return math.Min(float64(A[beginA]), float64(B[beginB]))
    }

    ia := int(math.Min(float64(m), float64(k/2)))
    ib := k - ia

    if A[beginA+ia-1] < B[beginB+ib-1] {
        return findKth(A, B, beginA+ia, beginB, k-ia)
    } else if A[beginA+ia-1] > B[beginB+ib-1] {
        return findKth(A, B, beginA, beginB+ib, k-ib)
    } else {
        return float64(A[beginA+ia-1])
    }
}

Reference