[LeetCode] Median of Two Sorted Arrays

There are two sorted arrays A and B 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)).

The trivial way, O(m + n):
Merge both arrays and the k-th smallest element could be accessed directly. Merging would require extra space of O(m+n). The linear run time is pretty good, but could we improve it even further?

A better way, O(k):
There is an improvement from the above method, thanks to readers who suggested this. (See comments below by Martin for an implementation). Using two pointers, you can traverse both arrays without actually merging them, thus without the extra space. Both pointers are initialized to point to head of A and B respectively, and the pointer that has the larger smaller (thanks to a reader for this correction) of the two is incremented one step. The k-th smallest is obtained by traversing a total of k steps. This algorithm is very similar to finding intersection of two sorted arrays.

The best solution, but non-trivial, O(lg m + lg n):
Although the above solution is an improvement both in run time and space complexity, it only works well for small values of k, and thus is still in linear run time. Could we improve the run time further?

The above logarithmic complexity gives us one important hint. Binary search is a great example of achieving logarithmic complexity by halving its search space in each iteration. Therefore, to achieve the complexity ofO(lg m  + lg n), we must halved the search space of A and B in each iteration.

We try to approach this tricky problem by comparing middle elements of A and B, which we identify as A_i and B_j. If A_i is between B_j and B_j-1, we have just found the i + j< + 1 smallest element. Why? Therefore, if we choose i and j such that i + j = k - 1, we are able to find the k-th smallest element. This is an important invariant that we must maintain for the correctness of this algorithm.

Summarizing the above,

Maintaining the invariant
    i + j = k - 1, 

If B_j-1 < A_i < B_j, then A_i must be the k-th smallest,
or else if A_i-1 < B_j < A_i, then B_j must be the k-th smallest.

If one of the above conditions are satisfied, we are done. If not, we will use i and j as the pivot index to subdivide the arrays. But how? Which portion should we discard? How about A_i and B_j itself?

We make an observation that when A_i < B_j, then it must be true that A_i < B_j-1. On the other hand, if B_j < A_i, then B_j < A_i-1. Why?

Using the above relationship, it becomes clear that when A_i < B_j, A_i and its lower portion could never be the k-th smallest element. So do B_j and its upper portion. Therefore, we could conveniently discard A_i with its lower portion and B_j with its upper portion.

If you are still not convince why the above argument is true, try drawing blocks representing elements in A and B. Try visualize inserting blocks of A up to A_i in front of B_j-1. You could easily see that no elements in the inserted blocks would ever be the k-th smallest. For the latter, you might want to keep the invariant i + j = k - 1 in mind to reason why B_j and its upper portion could never be the k-th smallest.

On the other hand, the case for A_i > B_j is just the other way around. Easy.

Below is the code and I have inserted lots of assertion (highly recommended programming style by the way) to help you understand the code. Note that the below code is an example of tail recursion, so you could technically convert it to an iterative method in a straightforward manner. However, I would leave it as it is, since this is how I derive the solution and it seemed more natural to be expressed in a recursive manner.

Another side note is regarding the choices of i and j. The below code would subdivide both arrays using its array sizes as weights. The reason is it might be able to guess the k-th element quicker (as long as the A and B is not differed in an extreme way; ie, all elements in A are smaller than B). If you are wondering, yes, you could choose i to be A's middle. In theory, you could choose any values for i and j as long as the invariant i+j = k-1 is satisfied.

//please refer to this link http://blog.csdn.net/jiyanfeng1/article/details/8619985

class Solution {
public:
    int find_kth_small(int A[], int m, int B[], int n, int k)
    {

        int i = (int)(((double) m)/(m+n)*(k-1));
        int j = k-1-i;
        // invariant: i + j = k-1  
        // Note: A[-1] = -INF and A[m] = +INF to maintain invariant  
        
        int Ai_1 = ((i == 0) ? INT_MIN :A[i-1]);
        int Bj_1 = ((j == 0) ? INT_MIN :B[j-1]);
        int Ai = ((i == m) ? INT_MAX :A[i]);
        int Bj = ((j == n) ? INT_MAX :B[j]);
        
        if(Ai == Bj)
            return Ai;
            
        if(Bj_1 < Ai && Ai < Bj)
            return Ai;
        if(Ai_1 < Bj && Bj < Ai)
            return Bj;

        // if none of the cases above, then it is either:    
        if(Ai < Bj)
        // exclude Ai and below portion  
        // exclude Bj and above portion 
            return find_kth_small(A+i+1, m-i-1, B, j, k-i-1);
        else /* Bj < Ai */
        // exclude Ai and above portion  
        // exclude Bj and below portion 
            return find_kth_small(A, i, B+j+1, n-j-1, k-j-1);
            //return find_kth_small(B+j+1, n-j-1, A, i, k-j-1);
    }
    double findMedianSortedArrays(int A[], int m, int B[], int n) {
        
        int cnt = m + n;
        
        if(cnt& 1 == 1)
        {
            return find_kth_small(A, m, B, n, cnt/2 +1 );
        }
        else 
            return (double)(find_kth_small(A, m, B, n, cnt/2 +1) + find_kth_small(A, m, B, n, cnt/2 ))/2.0;
    }
};    

http://blog.csdn.net/zxzxy1988/article/details/8587244

今天发现了leetcode上面一道题，觉得非常经典，记录之。

题目是这样的：给定两个已经排序好的数组（可能为空），找到两者所有元素中第k大的元素。另外一种更加具体的形式是，找到所有元素的中位数。本篇文章我们只讨论更加一般性的问题：如何找到两个数组中第k大的元素？不过，测试是用的两个数组的中位数的题目，Leetcode第4题 Median of Two Sorted Arrays
方案1：假设两个数组总共有n个元素，那么显然我们有用O(n)时间和O(n)空间的方法：用merge sort的思路排序，排序好的数组取出下标为k-1的元素就是我们需要的答案。
这个方法比较容易想到，但是有没有更好的方法呢？
方案2：我们可以发现，现在我们是不需要“排序”这么复杂的操作的，因为我们仅仅需要第k大的元素。我们可以用一个计数器，记录当前已经找到第m大的元素了。同时我们使用两个指针pA和pB，分别指向A和B数组的第一个元素。使用类似于merge sort的原理，如果数组A当前元素小，那么pA++，同时m++。如果数组B当前元素小，那么pB++，同时m++。最终当m等于k的时候，就得到了我们的答案——O(k)时间，O(1)空间。
但是，当k很接近于n的时候，这个方法还是很费时间的。当然，我们可以判断一下，如果k比n/2大的话，我们可以从最大的元素开始找。但是如果我们要找所有元素的中位数呢？时间还是O(n/2)=O(n)的。有没有更好的方案呢？
我们可以考虑从k入手。如果我们每次都能够剔除一个一定在第k大元素之前的元素，那么我们需要进行k次。但是如果每次我们都剔除一半呢？所以用这种类似于二分的思想，我们可以这样考虑：

Assume that the number of elements in A and B are both larger than k/2, and if we compare the k/2-th smallest element in A(i.e. A[k/2-1]) and the k-th smallest element in B(i.e. B[k/2 - 1]), there are three results:
(Becasue k can be odd or even number, so we assume k is even number here for simplicy. The following is also true when k is an odd number.)
A[k/2-1] = B[k/2-1]
A[k/2-1] > B[k/2-1]
A[k/2-1] < B[k/2-1]
if A[k/2-1] < B[k/2-1], that means all the elements from A[0] to A[k/2-1](i.e. the k/2 smallest elements in A) are in the range of k smallest elements in the union of A and B. Or, in the other word, A[k/2 - 1] can never be larger than the k-th smalleset element in the union of A and B.

Why?
We can use a proof by contradiction. Since A[k/2 - 1] is larger than the k-th smallest element in the union of A and B, then we assume it is the (k+1)-th smallest one. Since it is smaller than B[k/2 - 1], then B[k/2 - 1] should be at least the (k+2)-th smallest one. So there are at most (k/2-1) elements smaller than A[k/2-1] in A, and at most (k/2 - 1) elements smaller than A[k/2-1] in B.So the total number is k/2+k/2-2, which, no matter when k is odd or even, is surly smaller than k(since A[k/2-1] is the (k+1)-th smallest element). So A[k/2-1] can never larger than the k-th smallest element in the union of A and B if A[k/2-1]<B[k/2-1];
Since there is such an important conclusion, we can safely drop the first k/2 element in A, which are definitaly smaller than k-th element in the union of A and B. This is also true for the A[k/2-1] > B[k/2-1] condition, which we should drop the elements in B.
When A[k/2-1] = B[k/2-1], then we have found the k-th smallest element, that is the equal element, we can call it m. There are each (k/2-1) numbers smaller than m in A and B, so m must be the k-th smallest number. So we can call a function recursively, when A[k/2-1] < B[k/2-1], we drop the elements in A, else we drop the elements in B.


We should also consider the edge case, that is, when should we stop?
1. When A or B is empty, we return B[k-1]( or A[k-1]), respectively;
2. When k is 1(when A and B are both not empty), we return the smaller one of A[0] and B[0]
3. When A[k/2-1] = B[k/2-1], we should return one of them

In the code, we check if m is larger than n to garentee that the we always know the smaller array, for coding simplicy.

class Solution {
public:
    double findKth(int A[], int m, int B[], int n, int k)
    {
        //m is equal or smaller than n
        if (m > n)
            return findKth(B, n, A, m, k);
        if (m == 0)
            return B[k-1];
        if (k <= 1)
            return min(A[0], B[0]);

        int pa = min(k / 2, m), pb = k - pa;
        if (A[pa-1] < B[pb-1])
        {
            return findKth(A + pa, m - pa, B, n, k - pa);
        }
        else if(A[pa-1] > B[pb-1])
        {
            return findKth(A, m, B + pb, n - pb, k - pb);
        } else
            return A[pa-1];
    }

    double findMedianSortedArrays(int A[], int m, int B[], int n) {
        // Start typing your C/C++ solution below
        // DO NOT write int main() function
        int k = m + n;
        if (k & 0x1)
        {
            return findKth(A, m, B, n, k / 2 + 1);
        } else
        {
            return (findKth(A, m, B, n, k / 2) + findKth(A, m, B, n, k / 2 + 1)) / 2;
        }
    }
};