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

解答

求中位数需要根据数组长度是奇数还是偶数分别讨论，奇数长度时中位数为最中间的一个数，偶数长度时中位数为最中间的两个数的平均值，为了方便，可以实现一个比题目更一般化的函数，求A和B的第k小数的函数，那么中位数的问题很容易解决。

求一个有序数组的第k个数只需要O(1)的复杂度，现在有两个数组，显然花费额外空间以O(n)时间归并然后O(1)寻找不满足题目要求。

既然要求log时间复杂度，一般需要使用到二分思想。分别考虑A和B的第k/2个元素：

• 如果它们相等，则第k个数为其中的任意一个
• 如果A中的比较大，则B中前k/2个元素都不可能是第k个数了，因为这个数至少应该为A的第k/2个数，把B的前k/2去掉，然后重新寻找。
• 如果B中的比较大，则把A的前k/2个数去掉，重新寻找。

直到A和B中某个变为空时或者寻找第1个数时可以停止递归，直接找到结果。

注意，上面的k/2只是理想的简单情况，实际上A和B的长度可能不够k/2，或者k为奇数等，但这些不是主要问题，可以让A取第k/2个数字，然后A不够长，则取A的最后一个数字，然后B取剩下长度对应的那个数字，具体参考代码。

代码如下：

class Solution {
public:
    int f(int A[], int m, int B[], int n, int k) //寻找A和B的第k个数
    {
        if (m > n) return f(B, n, A, m, k); //保证A为较短的一个数组
        
        if (m == 0) return B[k-1];          //一个数组为空
        if (k == 1) return min(A[0], B[0]); //寻找第一个数
        
        int l1 = min(m, k / 2), l2 = k - l1;                          //A的一半位置，B的一半位置
        if (A[l1-1] == B[l2-1]) return A[l1-1];                       //两个数相等，找到答案
        else if (A[l1-1] < B[l2-1]) return f(A+l1, m-l1, B, n, k-l1); //B较大，去掉A的一半
        else return f(A, m, B+l2, n-l2, k-l2);                        //A较大，去掉B的一半
    }
    
    double findMedianSortedArrays(int A[], int m, int B[], int n) {
        int x = m + n;
        if (x & 1) //奇数长度中位数
            return f(A, m, B, n, x / 2 + 1);
        else       //偶数长度中位数
            return (f(A, m, B, n, x / 2) + f(A, m, B, n, x / 2 + 1)) / 2.0;
    }
};