[二分搜索|快速选择] leetcode 4 寻找两个正序数组的中位数

[二分搜索|快速选择] leetcode 4 寻找两个正序数组的中位数

1.题目

题目链接
给定两个大小为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。请你找出并返回这两个正序数组的中位数。
进阶： 你能设计一个时间复杂度为 O(log (m+n)) 的算法解决此问题吗？
示例1：

输入：nums1 = [1,3], nums2 = [2]
输出：2.00000
解释：合并数组 = [1,2,3] ，中位数 2

示例2：

输入：nums1 = [1,2], nums2 = [3,4]
输出：2.50000
解释：合并数组 = [1,2,3,4] ，中位数 (2 + 3) / 2 = 2.5

示例3：

输入：nums1 = [0,0], nums2 = [0,0]
输出：0.00000

示例4：

输入：nums1 = [], nums2 = [1]
输出：1.00000

示例5：

输入：nums1 = [2], nums2 = []
输出：2.00000

2.分析

2.1.直观思路

由于给出的两个数组是有序的，因此可以直接将两个数组合并形成一个新的有序数组。
这一过程与归并排序中“并”的过程十分相似：

double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
    vector<int> nums3;  //合并后的新数组
    int len1 = nums1.size();
    int len2 = nums2.size();
    int i = 0, j = 0;
    while(i < len1 || j < len2) { //遍历原两数组
        //如果数组1已经遍历完毕，将数组2的内容拷贝至新数组
        if(i > len1 - 1) {
            while(j < len2){
                nums3.push_back(nums2[j]);
                j++;
            }
            break;
        }
        //如果数组2已经遍历完毕，将数组1的内容拷贝至新数组
        if(j > len2 - 1) {
            while(i < len1){
                nums3.push_back(nums1[i]);
                i++;
            }
            break;
        }
		//选取数组1,2中较小的那个元素加入新数组
        if(nums1[i] <= nums2[j]) {
            nums3.push_back(nums1[i]);
            i++;
        }else{
            nums3.push_back(nums2[j]);
            j++;
        }
    }
    //取中位数
    if(nums3.size() % 2 != 0) {
        return nums3[nums3.size() / 2] * 1.0;
    }else{
        return (nums3[nums3.size() / 2] + nums3[nums3.size() / 2 - 1]) / 2.0;
    }
}

此时时间复杂度与空间复杂度均为O(m+n)。
实际上，由于两个原数组的长度是已知的，因此合并后新数组中位数的位置也是已知的，并不需要新数组num3，只需要一个变量index统计当前已经选择多少元素即可。这样可以将空间复杂度降为O(1)。

2.2.改进

题目要求的复杂度为O(log(m + n))，为对数级，似乎可以尝试使用二分的思想来解决。
为此，可以将“选取中位数”问题视为一个经典的选取数组中第K小元素的问题，即：

double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
    int len = nums1.size() + nums2.size();
    if (len % 2 != 0) {
        return getKthElement(nums1, nums2, (len + 1) / 2);
    } else {
        return (getKthElement(nums1, nums2, len / 2) + getKthElement(nums1, nums2, len / 2 + 1)) / 2.0;
    }
}

2.3.如何选取第K小元素

思路：令pivot1 = nums1[k / 2 - 1]，pivot2 = nums2[k / 2 - 1]。取pivot = min(pivot1, pivot2)。因此num1与nums2中比pivot小的元素最多有k / 2 - 1 + k / 2 - 1 = k - 2个，这样就可以确保pivot最多为第k - 1小的元素(甚至更小)。
如果pivot = pivot1，那么pivot及之前的数都不可能为第k小的元素(因为他们更小)，因此我们可以排除nums1[0~k/2-1]这些元素。
想象他们在数组中被删除，那么目前需要寻找的就是第“k - k / 2”小的元素了。这样就减小的问题的规模，直到最后“k = 1”，此时第一个元素便是我们要找的第K小元素。

3.代码

class Solution {
public:
    int getKthElement(const vector<int>& nums1, const vector<int>& nums2, int k) {
        int m = nums1.size();
        int n = nums2.size();
        //index之前的元素都是我们“已经排除掉的”
        int index1 = 0, index2 = 0;
        while (true) {
            //如果nums1中的元素全部排除了，那么我们选取nums2“排除后”的第k个元素即可
            if (index1 == m) {
                return nums2[index2 + k - 1];
            }
            //如果nums2中的元素全部排除了，那么我们选取nums1“排除后”的第k个元素即可
            if (index2 == n) {
                return nums1[index1 + k - 1];
            }
            //k=1则直接选取两数组中最小的那个元素返回即可
            if (k == 1) {
                return min(nums1[index1], nums2[index2]);
            }
            //选取pivot元素(注意越界)
            int newIndex1 = min(index1 + k / 2 - 1, m - 1);
            int newIndex2 = min(index2 + k / 2 - 1, n - 1);
            int pivot1 = nums1[newIndex1];
            int pivot2 = nums2[newIndex2];
            //nums1中的元素更小，那么k需要减去nums1中被排除的第index1个元素至第newIndex1个元素
            if (pivot1 <= pivot2) {
                k -= newIndex1 - index1 + 1;
                index1 = newIndex1 + 1; //index1之前的元素都被排除了
            }else { //反之同理
                k -= newIndex2 - index2 + 1;
                index2 = newIndex2 + 1;
            }
        }
    }

    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int len = nums1.size() + nums2.size();
        if (len % 2 == 1) {
            return getKthElement(nums1, nums2, (len + 1) / 2);
        } else {
            return (getKthElement(nums1, nums2, len / 2) + getKthElement(nums1, nums2, len / 2 + 1)) / 2.0;
        }
    }
};