leetcode 4. Median of Two Sorted Arrays

leetcode 4. Median of Two Sorted Arrays

Question

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

Example

Example 1:
nums1 = [1, 3]
nums2 = [2]

The median is 2.0

Example 2:
nums1 = [1, 2]
nums2 = [3, 4]

The median is (2 + 3)/2 = 2.5

问题

给定两个已排序的数组，找出其中的中位数。时间复杂度限制为O(log (m+n))。

分析

广义的问题可以定义为寻找p个排序数组中第k大的数。本题中p=2, k=n/2+1(n为奇数)或k=n/2，n/2+1(n为偶数)。

首先可以采用归并排序的思路，将两个数组先合并成一个大数组，然后再找出中位数。时间复杂度为O(m+n);
再优化一下，其实没有必要等合并成一个大数组之后再找，即找到对应位置的数字之后就停止合并，这样时间复杂度变成O(k)。但是由于k和m+n是线性的关系，时间复杂度还是O(m+n)。

题目要求O(log (m+n))，基于排序的log时间复杂度的算法，很快可以联想到二分查找。能不能利用二分查找的思路来做呢？二分查找的关键点在于，每次迭代都可以删除掉一部分不符合条件的数据（最多一半）。那我们能不能对多个排序数组同时进行快速删除呢？

这里要证明一个数学推理：
命题：给定两个排序数组a和b，设a的长度为m，b的长度为n。任取一个数r，r$\in$[1, m]。设k$\in$[1,m+n]，Nk为第k大的数。如果有a[r-1] < b[k-r-1]，那么a[0]-a[r-1]这r个数都比Nk要小。
证明：反证法，假如a[0]-a[r-1]这r个数不是都比Nk要小，即a[r-1] > Nk。由于a[r-1] < b[k-r-1]，则有b[k-r-1] > Nk。那么比Nk小的数最多有r-1 + k -r - 1 = k - 2个，矛盾。故命题得证。

如此一来我们就可以快速剔除大量不符合条件的数据了，可以让复杂度变成O(log (m+n))。

代码

class Solution {
public:
    int findKthNum(const vector<int> &a, int begA, const vector<int> &b, int begB, int k)
    {
        if (a.size() - begA > b.size() - begB)
            return findKthNum(b, begB, a, begA, k);
        if (a.size() - begA == 0)
            return b[begB + k - 1];
        if (k == 1)
            return min(a[begA], b[begB]);

        int pa = min(k / 2, (int)(a.size()) - begA);
        int pb = k - pa;

        if (a[begA + pa - 1] == b[begB + pb - 1])
            return a[begA + pa - 1];
        else if (a[begA + pa - 1] < b[begB + pb - 1])
            return findKthNum(a, begA + pa, b, begB, k - pa);
        else 
            return findKthNum(a, begA, b, begB + pb, k - pb);
    }

    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int n = nums1.size() + nums2.size();

        if (n & 0x1)
        {
            return findKthNum(nums1, 0, nums2, 0, n / 2 + 1);
        }
        else
        {
            return (findKthNum(nums1, 0, nums2, 0, n / 2) + findKthNum(nums1, 0, nums2, 0, n / 2 + 1)) / 2.0;
        }
    }
};

总结

本题较难，需要从一个数学结论出发，进而快速剔除不符合条件的数据。值得注意的是，findKthNum函数有三个退出条件，分别是：
1、a数组为空，直接返回b的第k个数；
2、a数组和b数组在对应位置上的数相等，直接返回这个数；
3、k等于1时，直接a数组和b数组第一个数较小的那个数。

其中第三个条件是不可或缺的，不然有些情况会数组越界（可以举个简单的例子试试）。