【LeetCode】4. Median of Two Sorted Arrays (2 solutions)

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)).

解法一：保底做法，O(m+n)复杂度

按照归并排序的思路，先归并，再找中间值。

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int m = nums1.size();
        int n = nums2.size();
        vector<int> array = merge(nums1, nums2);
        return ((double)(array[(m+n-1)/2]+array[(m+n)/2]))/2;
    }
    vector<int> merge(vector<int> A, vector<int> B)
    {
        vector<int> ret;
        int m = A.size();
        int n = B.size();
        int i = 0;
        int j = 0;
        while(i < m && j < n)
        {
            if(A[i] <= B[j])
            {
                ret.push_back(A[i]);
                i ++;
            }
            else
            {
                ret.push_back(B[j]);
                j ++;
            }
        }
        if(i == m)
        {
            while(j < n)
            {
                ret.push_back(B[j]);
                j ++;
            }
        }
        if(j == n)
        {
            while(i < m)
            {
                ret.push_back(A[i]);
                i ++;
            }
        }
        return ret;
    }
};

解法二：类似二分查找，复杂度O(log(m+n))

这个解法大概思路很简单，就是A数组的中间元素与B数组的中间元素相比较，从而删掉较小元素所在数组的前一半和较大元素所在数组的后一半。递归下去。

其实要正确实现以上解法还是比较tricky的，因为需要涉及到很多边界情况，我们来一一解释：

第一步，我们需要将题目改为寻找第k小的元素findKthSortedArrays。（这里k从1开始计算）

原因在于，如果我们执着于median，那么在删除一半左右元素形成的子问题中，很难保证仍然是寻找median。可能变为寻找median前一个或者后一个。

（如果仔细思考过这个题就能理解这句话）

改成第k小元素的好处就是，我们可以将删掉的元素考虑进来，在子问题中不断改变k的值。

（例如：本来我们需要寻找的median是第5个数，删掉前面2个数之后，在子问题中就变为寻找第5-2=3个数）

考虑A、B数组总数的奇偶性，就转化为调用findKthSortedArrays的问题了。

第二部：实现findKthSortedArrays

（一）首先，我们规定A数组比B数组短。

这样处理的好处在于：我们所需的(k-1)/2位置可能大于某个数组总长度，规定A短之后，只需要考虑超过A的长度，不需要再复制代码分情况讨论了。

这里需要斟酌一下：为什么不是k/2？ k/2±1？而是(k-1)/2

我的思考如下：

如果k==1,2，就是需要比较头两个元素，因此下标为0

如果k==3,4，就是需要比较第1个元素，因此下标为1

综上归纳而得。

（二）其次，把特殊情况处理掉。

(1)A数组为空时，即返回B数组第k个元素。

(2)k==1时，即返回A[0]、B[0]中小的那个

（三）接下来再分为两种情况：

(k-1)/2位置是否超过A数组长度？

（1）若超过，则A数组派出的代表Acandi就是最后元素A[m-1]，B派出的代表Bcandi是B[k-m-1]

(a)Acandi==Bcandi，那么正好有k-2个元素比Acandi、Bcandi小，所以第k个元素就是Acandi/Bcandi

(b)Acandi > Bcandi，那么最多只有m-1+k-m-1=k-2个元素比Bcandi小，因此包括Bcandi在内之前的k-m个B数组元素肯定不是第k个数，所以删去，子问题变为寻找第k-(k-m)个元素

(c)Acandi < Bcandi，那么最多只有m-1+k-m-1=k-2个元素比Acandi小，因此包括Acandi在内之前的整个A数组元素m个元素肯定不是第k个数，所以删去，子问题变为寻找第k-m个元素

（2）若不超过，则A数组派出的代表Acandi就是A[(k-1)/2]，B派出的代表Bcandi是B[(k-1)/2]

(a)Acandi==Bcandi，那么正好有(k-1)/2+(k-1)/2=k-1个元素比Acandi、Bcandi小，所以第k个元素就是Acandi/Bcandi

如果不相等，对于Acandi 、Bcandi本身是否需要取舍就要注意分析了。

经过举几个例子简单分析就很容易发现，k为奇数时，需要保留小的candidate，舍弃大的。

而k为偶数时，需要保留大的candidate，舍弃小的。

(b)Acandi > Bcandi

(b.1)k为奇数，保留Bcandi，删除Bcandi前面的(k-1)/2个元素，删除A及A后面的元素（保留A中前(k-1)/2个元素）

(b.2)k为偶数，保留Acandi，删除Bcandi及前面的(k-1)/2+1个元素，删除A后面的元素（保留A中前(k-1)/2+1个元素）

(c)Acandi < Bcandi

同(b)，略去不赘述了。

class Solution {
public:
    double findKthSortedArrays(int A[], int m, int B[], int n, int k)
    {//k starts from 1
    
        //make sure A is shorter than B
        if(m > n)
            return findKthSortedArrays(B, n, A, m, k);
    
        //special case1: A empty
        if(m == 0)
            return B[k-1];
        //special case2: k==1 (m>0 is guaranteed)
        if(k == 1)
            return min(A[0], B[0]);
        
        int Acandi, Bcandi;
        if((k-1)/2 >= m)
        {//A[(k-1)/2] out of range
            Acandi = A[m-1];
            Bcandi = B[k-m-1];
            if(Acandi == Bcandi)
                return Acandi;
            else if(Acandi > Bcandi) 
            //for A: no skip
            //for B: skip the k-m smaller elements (including Bcandi)
                return findKthSortedArrays(A, m, B+k-m, n-(k-m), k-(k-m));
            else
            //for A: skip the m   smaller elements
            //for B: skip the k-m larger  elements
                return findKthSortedArrays(A+m, 0, B, n-(k-m), k-m);
        }
        else
        {
            //1,2->index0; 3,4->index1; ...
            Acandi = A[(k-1)/2];
            Bcandi = B[(k-1)/2];
            if(Acandi == Bcandi)
                return Acandi;
            else if(Acandi > Bcandi) 
            {
            //for A: skip the larger  elements
            //for B: skip the smaller elements
                if(k%2 == 1)
                //keep the smaller candidate, skip the larger
                    return findKthSortedArrays(A, (k-1)/2, B+(k-1)/2, n-(k-1)/2, k-(k-1)/2);
                else
                //keep the larger candidate, skip the smaller
                    return findKthSortedArrays(A, (k-1)/2+1, B+(k-1)/2+1, n-((k-1)/2+1), k-((k-1)/2+1));
            }
            else
            {
            //for A: skip the smaller elements
            //for B: skip the larger  elements
                if(k%2 == 1)
                //keep the smaller candidate, skip the larger
                    return findKthSortedArrays(A+(k-1)/2, m-(k-1)/2, B, (k-1)/2, k-(k-1)/2);
                else
                //keep the larger candidate, skip the smaller
                    return findKthSortedArrays(A+(k-1)/2+1, m-((k-1)/2+1), B, (k-1)/2+1, k-((k-1)/2+1));
            }
        }
    }
    double findMedianSortedArrays(int A[], int m, int B[], int n) 
    {
        if((m+n)%2 == 0)
        {//average of two medians
            return (findKthSortedArrays(A, m, B, n, (m+n)/2) + findKthSortedArrays(A, m, B, n, (m+n)/2+1))/2;
        }
        else
        {
            return findKthSortedArrays(A, m, B, n, (m+n)/2+1);
        }
    }
};