Leetcode Hot100不熟练题目 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

提示：

nums1.length == m
nums2.length == n
0 <= m <= 1000
0 <= n <= 1000
1 <= m + n <= 2000
-106 <= nums1[i], nums2[i] <= 10

2、思路

1. 问题转换

对于给定的两个数组，其长度是固定的m、n,所以我们寻找的中位数可以这样分析。

• 当m + n为奇数的时候，我们需要得到两数组中第**（m + n + 1) / 2**小的数值。
• 当m + n为偶数的时候，我们计算的就是第**（m + n) / 2** 小的数和第**（m + n) / 2 + 1**小的数的平均值。

所以我们的目标就变成了，找到这两个数组中第k小的元素。

2. 找到两数组中第k小的元素

对于log的时间复杂度，我们首先就想到了二分的方法。设第一个数组为a,第二个数组为b。

1. 一般的情况

我们取a的第k / 2个数和b的k / 2个数进行比较。并假设a[k / 2] <= b[k / 2]。
这种情况下，对于a[k/2]来说，最多只有a的前 k / 2 - 1个元素和b的前 k / 2 - 1个元素小于它，一共 k / 2 - 1 + k / 2 - 1= k - 2个元素。
所以a[k / 2]一定不是第k小的数，它最多是第k - 1小的元素。因此a的第1个到第k / 2个元素，都不可能是正确答案，将他们排除。也就是将k减去k / 2，也就是我们要在剩下的元素中找到第 k - k / 2 小的元素。

2.特殊情况

1. k / 2 大于 某个数组剩余元素的数量。此时我们选择最后一个元素作为比较对象，并在排除时如实地排除相应数量的元素(该数组中剩余元素的数量，而不是k / 2)
2. 当某个数组为空时，可以直接返回另一个数组的第k个元素
3. 当k为1时，可以直接返回两个数组的第一个元素中的最小值

3、代码

class Solution {
public:
    // 2. 转换的目标，求第k个元素
    double findk(vector<int> &nums1, vector<int> & nums2,int k)
    {
        //idx1代表nums1的有效起点，idx2代表nums2的有效起点
        int idx1 = 0, idx2 = 0;
        int n = nums1.size(), m = nums2.size();
        while(true)
        {
            //3. 特殊情况处理
            //3. 1 某个数组元素数量为0,返回另一个的第k个元素
            if(idx1 == n) return nums2[idx2 + k - 1];
            if(idx2 == m) return nums1[idx1 + k - 1];
            //3. 2 k == 1
            if(k == 1) return min(nums1[idx1],nums2[idx2]);
            //3. 3 其中一个数组元素数量不足k / 2
            int newidx1 = min(idx1 + k / 2 - 1, n - 1);
            int newidx2 = min(idx2 + k / 2 - 1, m - 1);
            if(nums1[newidx1] <= nums2[newidx2])
            {
                // newidx1 - idx1 的数量代表newidx1之前的元素的数量，它本身也要被排除，所以要+1
                k -= newidx1 - idx1 + 1;
                // idx1 是新的有效起点，是newidx1的下一个，所以+1
                idx1 = newidx1 + 1;
            }else{
                k -= newidx2 - idx2 + 1;
                idx2 = newidx2 + 1;
            }
        }
        return 0.0;
    }
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int total = nums1.size() + nums2.size();
        //1. 对于奇偶不同的处理
        // 奇数的中位数是total + 1 / 2
        // 偶数分别是计算 total / 2 和 total / 2 + 1
        if(total % 2) return findk(nums1,nums2,(total + 1) / 2);
        else return (findk(nums1,nums2,total / 2) + findk(nums1,nums2,total / 2 + 1)) / 2;
    }
};

4、参考资料

官方题解