题目

给定两个大小分别为mn的正序(从小到大)数组nums1nums2。请你找出并返回这两个正序数组的中位数 。

算法的时间复杂度应该为O(log(m+n))

示例 1:

  1. 输入:nums1 = [1,3], nums2 = [2]
  2. 输出:2.00000
  3. 解释:合并数组 = [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
  • -10^6 <= nums1[i], nums2[i] <= 10^6

解题方案

二分法

设两数组分别长为mn。当m+n为偶数是,两数组中位数为第(m+n)/2和第(m+n)/2+1小个元素和的一半;当m+n为奇数时,中位数为两数组中第(m+n)/2+1小个元素。因此,问题转化为求两数组中第k小个元素。

采用二分法,比较两数组第k/2-1个元素,存在以下情况:

  • A[k/2-1]≤B[k/2-1]。对于A[k/2-1],前面最多有k-2个元素比他小,因此A[0]~A[k/2-1]中不包含第k小的数,可以舍弃。问题进一步转换为求解A[k/2: ]B[0: ]中第k-k/2小个数。
  • A[k/2-1]>B[k/2-1]。对于B[k/2-1],前面最多有k-2个元素比他小,因此B[0]~B[k/2-1]中不包含第k小的数,可以舍弃。问题进一步转换为求解B[k/2: ]A[0: ]中第k-k/2小个数。

存在以下特殊情况:

  • [k/2-1]越界,则选取对应数组最后一个数即可。
  • k=1,则直接比较两数组首元素。
  • 若一个数组已经全部排除,直接在剩下的数组中选取第k小的数即可。

该方法时间复杂度为O(log(m+n))

C++代码

class Solution {
public:
    int kth_min(vector<int>& nums1, vector<int>& nums2, int k) {
        int m = nums1.size();
        int n = nums2.size();

        int offset1 = 0;
        int offset2 = 0;

        while(true) {
            if(offset1==m)  return nums2[offset2 + k - 1];
            if(offset2==n)  return nums1[offset1 + k - 1];
            if(k==1)    return min(nums1[offset1], nums2[offset2]);

            int new_offset1 = min(offset1 + k / 2 - 1, m - 1);
            int new_offset2 = min(offset2 + k / 2 - 1, n - 1);
            if(nums1[new_offset1] <= nums2[new_offset2]) {
                k = k - (new_offset1 - offset1 + 1);
                offset1 = new_offset1 + 1;
            }
            else {
                k = k - (new_offset2 - offset2 + 1);
                offset2 = new_offset2 + 1;
            }
        }
    }

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

分割数组

中位数将有序数组分割成了前后元素数量相等的两个子数组,从该定义出发可以更快的找到中位数。对两个数组AB,长度分别为mn。设在A数组的i处分割,B数组的j处分割。为了保证分割后左侧数组的集合中元素数量与右侧数组集合中的相等,有如下等式。
4. 寻找两个正序数组的中位数 - 图1
式中为相除取整操作。当i0增大时,恒有A[i-1]<=B[j]
因此,对元素数较小的数组二分查找满足上述不等式的最大下标,即可求得中位数。
4. 寻找两个正序数组的中位数 - 图2
该方法时间复杂度为O(log(min(m, n)))

C++代码

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int m = nums1.size();
        int n = nums2.size();
        if (m>n) return findMedianSortedArrays(nums2, nums1);

        int left = 0;
        int right = m;
        int left_max = 0;
        int right_min = 0;

        while (left<=right) {
            int i = (left + right) / 2;
            int j = (m + n + 1) / 2 - i;

            int nums1_im1 = i==0 ? INT_MIN : nums1[i-1];
            int nums1_i = i==m ? INT_MAX : nums1[i];
            int nums2_jm1 = j==0 ? INT_MIN : nums2[j-1];
            int nums2_j = j==n ? INT_MAX : nums2[j];

            if(nums1_im1 <= nums2_j) {
                left_max = max(nums1_im1, nums2_jm1);
                right_min = min(nums1_i, nums2_j);
                left = i + 1;
            }
            else {
                right = i - 1;
            }
        }

        return (m+n)%2==0 ? (left_max + right_min) / 2.0 : left_max;
    }
};