【LeetCode-二分查找】寻找两个正序数组的中位数

题目描述

给定两个大小为 m 和 n 的正序(从小到大)数组 nums1 和 nums2。

请你找出这两个正序数组的中位数,并且要求算法的时间复杂度为 O(log(m + n))。

你可以假设 nums1 和 nums2 不会同时为空。

示例:

nums1 = [1, 3]
nums2 = [2]
则中位数是 2.0

nums1 = [1, 2]
nums2 = [3, 4]
则中位数是 (2 + 3)/2 = 2.5

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

思路

因为题目要求算法的时间复杂度为 O(log(m + n)),log 的时间复杂度一般是使用二分查找来做。

因为两个数组是有序的,所以我们可以一半一半地对数组进行寻找。假设两个数组如下

图来自这篇题解

假设我们要找第 k = 7 小的数字,则我们可以比较 nums1[k/2] 和 nums2[k/2] 之间的大小,如果 nums2[k/2]<nums1[k/2],因为数组是升序的,则 nums2[k/2] 之前的数字都不会是第 k 小的数字,可以直接删除

黄色部分表示已删除。此时我们已经删除了 7/2 = 3 个数字,所以接下来我们要找第 k = 7 - 3 = 4 小的数字。同样地,我们继续比较 nums1[k/2] 和 nums2[k/2] 之间的大小,根据大小关系判断如何删除元素。

这个过程会有一个问题,因为我们每次都是比较两个数组的第 k/2 个元素,但可能会出现数组长度小于 k/2 的情况,如下

解决办法就是我们取 min(k/2, len) - 1 对应的元素进行比较,len 为数组的剩余长度。

代码如下:

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int m = nums1.size();
        int n = nums2.size();
        int left = (m+n+1)/2;
        int right = (m+n+2)/2;
        return (find(nums1, 0, m-1, nums2, 0, n-1, left) + find(nums1, 0, m-1, nums2, 0, n-1, right)) / 2;
    }

    double find(vector<int>& nums1, int left1, int right1, vector<int>& nums2, int left2, int right2, int k){
        int len1 = right1-left1+1;
        int len2 = right2-left2+1;
        if(len1>len2) return find(nums2, left2, right2, nums1, left1, right1, k); // 保证 len1 不大于 len2
        if(len1==0) return nums2[left2+k-1];  // 如果 len1==0,直接返回当前 nums2 的第 k 个元素
        if(k==1) return min(nums1[left1], nums2[left2]);  // 如果 k==1,则返回两个数组头更小的那个元素

        int i = left1 + min(k/2, len1)-1;  // 第 k/2 个元素,为了防止数组长度小于 k/2,用了 min 函数 
        int j = left2 + min(k/2, len2)-1;  // 同上
        if(nums1[i]<nums2[j])
            return find(nums1, i+1, right1, nums2, left2, right2, k-(i-left1+1)); // 别忘了最后一个参数
        else 
            return find(nums1, left1, right1, nums2, j+1, right2, k-(j-left2+1));
    }
};

代码中还有一点,就是我们同时计算了第 left 小的元素和第 right 小的元素,然后求两者的平均值作为答案。这样做的原因是,假设两个数组长度分别为 len1 和 len2,如果 len1+len2 为偶数,则我们需要求中间两个数的平均值,中间两个数就是第 left 小的元素和第 right 小的元素;如果 len1+len2 为奇数,则中间的那个数就是中位数,这时第 left 小的元素和第 right 小的元素是一样的,会发生重复计算。

参考

https://leetcode-cn.com/problems/median-of-two-sorted-arrays/solution/xiang-xi-tong-su-de-si-lu-fen-xi-duo-jie-fa-by-w-2/

posted @ 2020-09-12 10:31  Flix  阅读(1006)  评论(0编辑  收藏  举报