leetcode hot100 之 寻找两个正序数组的中位数

题目

给定两个正序数组，寻找这两个数组的中位数。

示例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

原题链接： https://leetcode-cn.com/problems/median-of-two-sorted-arrays/

思路

思路1

假设数组长度分别为m 和 n，k = (m+n)/2。如果 m+n 为奇数，则中位数第 k+1 个数；如果 m+n 为偶数，则中位数是第 k 个数 和第 k+1 个数的平均值。
对于两个数组，我们可以分别维护一个指针，每一次对比当前指针下两个数组的值，从而可以确认第 x 小的值是哪一个（这里 x = i+j，其中 i 和 j 分别为指针下标），直到找到第 k+1 个数为止。
同时，针对偶数的情况，我们需要维护两个变量来记录第 x-1 小的值和第 x 小的值。

• 复杂度分析
  • 时间复杂度 O(m+n)。
  • 空间复杂度 O(1)。

思路2（高阶）

如果要求时间复杂度为 O(log(m+n))，那就只有考虑二分法了。
根据思路1，我们的目标是找到第 k+1 个数。那么可以通过二分法，先找出 top(k+1) 小的数。为了方便表示，设 K=k+1。二分的过程为，每次先找出 topK 小当中的 K/2 个数，为了保证这 K/2 个数一定在 topK 小的数当中，我们只需要对比两个数组当中第 K/2 个数，即 nums1[K/2-1] 和 nums2[K/2-1] 这两个数的值，如果 nums2[K/2-1] 小于或等于 nums1[K/2-1]， 则 nums2[0] ~ nums2[K/2-1] 这 K/2 个数一定是在 topK 小的数当中的。排除掉 K/2 个数之后，就进一步找剩下的 X=(K - K/2) 个数，直到找完 topK 为止。
如果某个数组遍历完成，或者某个数组为空的情况，可以对另一个数组直接移动X-1步，提前结束二分。
二分的过程中，走到 X = 1 这一步的时候，我们已经取得了 k = K-1 的集合，对比两个数组头部的数字，取最大的一个即是我们想要的第 k 个数。然后再找出 X = 1 这一步的数字，即是第 k+1 个数。

• 复杂度分析
  • 时间复杂度 O(log(m+n))。
  • 空间复杂度 O(1)。

代码

代码1

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int k = (nums1.size() + nums2.size()) / 2;
        double pre, cur;
        int i = 0;
        int j = 0;
        while(i + j != k + 1) {
            pre = cur;
            if (i < nums1.size() && j < nums2.size()) {
                if (nums1[i] > nums2[j]) {
                    cur = nums2[j];
                    j++;
                }
                else {
                    cur = nums1[i];
                    i++;
                }
            }
            else if (i < nums1.size()) {
                cur = nums1[i];
                i++;
            }
            else if (j < nums2.size()) {
                cur = nums2[j];
                j++;
            }
        }
        if ( (nums1.size() + nums2.size()) % 2 == 0) {
            return (pre + cur) / 2;
        }
        else {
            return cur;
        }
    }
};

代码2

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int index1 = -1;
        int index2 = -1;
        double pre, cur;
        int k = (nums1.size() + nums2.size()) / 2;
        cur = findTopK(nums1, index1, nums2, index2, k+1, pre);
        if ((nums1.size() + nums2.size()) % 2 == 1) {
            return cur;
        }
        else {
            return (cur + pre) / 2;
        }
    }
    double findTopK(vector<int>& nums1, int& index1,
                    vector<int>& nums2, int& index2,
                    int k, double& pre) {
        int mid = k/2;
        int step1 = mid;
        int step2 = mid;
        
        // 当 k == 1 时，则进行最后一轮比较和查找，同时找出倒数第二大的数
        if (k == 1) {
            int pre1, pre2;
            pre1 = index1 == -1 ? INT_MIN: nums1[index1];
            pre2 = index2 == -1 ? INT_MIN: nums2[index2];
            pre = max(pre1, pre2);
            if (index1 == nums1.size() - 1) {
                return nums2[index2 + 1];
            }
            if (index2 == nums2.size() - 1) {
                return nums1[index1 + 1];
            }
            if (nums1[index1 + 1] > nums2[index2 + 1]) {
                return nums2[index2 + 1];
            }
            else {
                return nums1[index1 + 1];
            }
        }

        // 如果某个数组已经遍历完，则直接对另一个数组移动 k-1步; 同时也覆盖数组为空的情况
        if (index1 == nums1.size() - 1) {
            index2 += k - 1;
            return findTopK(nums1, index1, nums2, index2, 1, pre);
        }
        if (index2 == nums2.size() - 1) {
            index1 += k - 1;
            return findTopK(nums1, index1, nums2, index2, 1, pre);
        }
        // 如果某个数组移动 mid 步后已越出数组，则调小step，只移到数组最后一个
        if (index1 + step1 > nums1.size() - 1) {
            step1 = nums1.size() - 1 - index1;
        }
        if (index2 + step2 > nums2.size() - 1) {
            step2 = nums2.size() - 1 - index2;
        }
        
        // 比较两个数组第 index+step 位，找出topK 中的数
        if (nums1[index1 + step1] > nums2[index2 + step2]) {
            index2 += step2;
            return findTopK(nums1, index1, nums2, index2, k - step2, pre);
        }
        else {
            index1 += step1;
            return findTopK(nums1, index1, nums2, index2, k - step1, pre);
        }
    }
};