[LeetCode HARD] 4 Median of Two Sorted Arrays

LeetCode has become the de facto go-to website for tech interviews. As a former competitive programmer, the LeetCode “hard”, “medium”, and “easy” problems are more like “easy”, “very easy”, “extremely easy” for me. I think it would be a good idea for me to put my solutions of some of the LeetCode hard problems here. Hope it would be helpful for people who are preparing for tech interviews.

This is my first LeetCode solution article, for LeetCode Problem 4, Medium of Two Sorted Arrays.

Problem

The problem is simple: given two sorted integer arrays, find the medium element of the elements from the two arrays combined.

Solution

First, instead of finding the medium of the elements, we can solve a generalized version of this problem: giving two sorted arrays, find the k-th smallest element for any k.

Now assume that k-th element is smaller or equal to the medium. Consider the two medium elements of both arrays, m_1 and m_2, if m_1 >= m_2, we know that any element e after m_1 in array 1 would be >= m_1 >= m_2. Thus all those elements would be larger than all elements before m_1 in array 1 and all elements before m_2 in the second array, which are half of all elements. Thus e is larger than the medium and can’t be the solution. Thus we can eliminate all element e, i.e., all elements after m_1 in array 1.

Similarly, we can eliminate half of the elements in array 2 if m_1 <= m_2. And similarly, if the k-th element is larger than the medium, we can eliminate the first half of array 1 or array 2.

In every iteration, we can either eliminate half of the elements in array 1 or eliminate half of the elements in array 2. Therefore, in O(log(m)+log(n)) = O(log(m+n)) iteration, we can reduce one of the arrays to length 1, in which case we can fine the solution in O(1) time.

Note

The algorithm and the proofs above are not rigid. I prefer to focus on the core idea of the algorithm instead of wasting time on the details which would make the this article incomprehensible. You can see the source code in the end if you are interested in implementation details.

Takeaway

There are two takeaways of solving this problem. The first one is finding a generalized problem so that we can find a recursive solution. It’s hard to find a recursive relation between finding the medium of two longer arrays and finding the medium of two shorter arrays. But it’s easier to find the relationship between finding the k-th element of two longer arrays and the k’-th element between two shorter arrays.

The second one is how to find the recursive step. If we can reduce the size of the problem to any fraction of the original size in constant time, we can find a O(log(n)) solution. We don’t have to reduce the size of the problem by half. In this particular problem, we only reduce half of the elements in one of the two arrays, about 1/4 of the total elements. But this algorithm is still O(log(n)).

Source Code

My solution is posted below. I write it in good old C because it’s easier to pass sub-arrays, and I am confident that getting a sub-array in C is O(1) without looking up the documents.

int findElement(int* nums1, int nums1Size, int* nums2, int nums2Size, int k) {
    // We make sure the first array is longer than the second array 
    // to make our lives easier later.
    if (nums1Size > nums2Size)
        return findElement(nums2, nums2Size, nums1, nums1Size, k);
    
    // base cases
    if (nums1Size == 0)
        return nums2[k];
    if (nums1Size == 1) {
        int e = nums1[0];
        if (k == 0)
            return (e < nums2[0]) ? e : nums2[0]; if (k >= nums2Size) 
            return (e > nums2[nums2Size - 1]) ? e : nums2[nums2Size - 1];
        if (e >= nums2[k]) {
            return nums2[k];
        }
        return (e >= nums2[k-1]) ? e : nums2[k-1];
    }
    
    // induction step
    if (k + k <= nums1Size + nums2Size) {
        if (nums1[nums1Size/2] < nums2[nums2Size/2])
            return findElement(nums1, nums1Size, nums2, (nums2Size+1) / 2, k);
        else
            return findElement(nums1, (nums1Size+1) / 2, nums2, nums2Size, k);
    }
    else {
        if ( nums1[(nums1Size-1) / 2] < nums2[(nums2Size-1) / 2] ) 
            return findElement(&nums1[nums1Size/2], (nums1Size+1) / 2, 
                nums2, nums2Size, k - nums1Size/2);
        else
            return findElement(nums1, nums1Size, 
                &nums2[nums2Size/2], (nums2Size+1) / 2, k - nums2Size/2);
    }
}

double findMedianSortedArrays(int* nums1, int nums1Size, int* nums2, int nums2Size) {
    // if there are even number of elements, we need to find 
    // the average of the two medium elements.
    int total = nums1Size + nums2Size;
    if (total & 1)
        return (double)findElement(nums1, nums1Size, nums2, nums2Size, total/2 );
    else 
        return ((double)findElement(nums1, nums1Size, nums2, nums2Size, total/2-1 ) + 
            (double)findElement(nums1, nums1Size, nums2, nums2Size, total/2) ) / 2.;
}