[LeetCode] Median of Two Sorted Arrays

Problem Statement (link):
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)).
Analysis:

- An O(m+n) solution
It is obvious we can merge the two array into one big array in O(m+n) time, and find the median in O(1). In addition, this requires O(m+n) space.

- An O(k) solution
As the two arrays are sorted and we know the index of median of the merged array - either (m+n)/2 in odd case or ((m+n)/2+(m+n)/2-1)/2 in even case, we could simply count from the start of two arrays until we reach the index.

- An O(log (m + n)) solution
The above two algorithms are in linear time, as we are required to solve the program in log time, it's only logical that we consider binary search. In addition, this problem can be generalized to finding the k-th smallest element in two sorted arrays, hereby k = (m+n)/2+1, note that k is not the index that starts from 0, instead it starts from 1.

In each recursion, we find the medians of two arrays. The two medians will split the two arrays into four parts as shown below:

Figure, The medians separates two arrays into four parts. Courtesy of Lei Zhang.

As shown in the picture, we need to consider the two conditions in each round and discard one part as we do in binary search. The logic goes like this:

- k >= m/2+n/2+1 && am/2 > bn/2, discard Section 3;
- k >= m/2+n/2+1 && am/2 < bn/2, discard Section 1;
- k < m/2+n/2+1 && am/2 > bn/2, discard Section 2;
- k < m/2+n/2+1 && am/2 < bn/2, discard Section 4;

For instance, if k >= m/2+n/2+1, we know that the k-th item must be in larger half of the merged array; further, if  am/2 > bn/2, then we could discard section 3 as its elements are too small to be median. The other three conditions follow the same logic.

Code:

class Solution {
public:
    // generalized to finding kth smallest element problem
    double findMedianSortedArrays(int A[], int m, int B[], int n) {
        int total=m+n;
        if (total%2==1)
            return recur(A,m,B,n,total/2+1);
        else
            return (recur(A,m,B,n,total/2)+recur(A,m,B,n,total/2+1))/2.0;
    }

    // find k-th smallest element
    double recur(int A[], int m, int B[], int n, int k) {
        if (m<=0) return B[k-1];
        if (n<=0) return A[k-1];
        if (k<=1) return min(A[0], B[0]);

        int mid=m/2+n/2+1;
        if (B[n/2]>=A[m/2]) {
            if (mid>=k)
                return recur(A, m, B, n/2, k);
            else
                return recur(A+m/2+1, m-(m/2+1), B, n, k-(m/2+1)); 
        }
        else {
            if (mid>=k)
                return recur(A, m/2, B, n, k);
            else
                return recur(A, m, B+n/2+1, n-(n/2+1), k-(n/2+1));
        }
    }
};

Takeaways:
Sometimes we cannot simplify integer calculation argument, e.g., m/2-1 != m-(m/2+1)