【LeetCode Hot 100】4. 寻找两个正序数组的中位数

题目描述

要求出两个数组的中位数，第一想法当然是将这两个数组进行归并排序，然后直接得到排序后长数组的中位数。由于本题的两个数组都是排序后的数组，因此省去了排序的步骤。这种方法的时间复杂度为$O(m+n)$，空间复杂度由于要存储排序后的长数组，所以也是$O(m+n)$。

有没有相对更简单的方法呢？由于我们只需要求出中位数，并不需要得到合并后的数组，所以我们可以从左往右逐个元素遍历并计数，只要我们达到了中位数所需要的下标，就可以停下来并计算中位数了。需要注意的是，如果总长度为奇数，中位数应该是中间两个元素的算术平均，因此除了当前元素之外，还需要维护前一个元素。显然，这种方法的时间复杂度为$O(\frac{m+n}{2})$，空间复杂度则相对第一种方法优化为了常数级。

// C++
class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int m = nums1.size(), n = nums2.size();
        int len = m + n;
        int prev = -1, curr = -1;
        int i = 0, j = 0;
        for (int cnt = 0; cnt <= len / 2; cnt++) {
            prev = curr;
            if (i < m && j < n) {
                if (nums1[i] < nums2[j]) {
                    curr = nums1[i++];
                } else {
                    curr = nums2[j++];
                }
                continue;
            }
            if (i < m) {
                curr = nums1[i++];
            }
            if (j < n) {
                curr = nums2[j++];
            }
        }
        if (len % 2) {
            return curr;
        }
        return (prev + curr) / 2.0;
    }
};

但是，题目描述中要求解法的时间复杂度为对数级别，看到这个要求，就要想到与二分有关。

要找到两个数组的中位数，就是要找到第$\frac{m+n}{2}$小（和第$\frac{m+n}{2}+1$小）的元素。这个方法则给出了更通用的方法，用于求出第$k$小的元素。

要找到两个升序数组总共的第$k$小的元素，我们可以先比较两个数组下标为k/2 - 1的元素，假设a[k/2 - 1] < b[k/2 - 1]，那么a数组的前面几个元素只会更小，我们即使假设b数组的前面几个元素都比a[k/2 - 1]小，此时该元素也“仅仅”是总共第$\frac{k}{2}-1$小的元素，也就是说它最大也只能达到这个程度，而绝不可能是第$k$小的元素，那么我们就可以放心地把a[0..k/2-1]这个子数组全部“丢弃”不看，这样我们一次性排除了$\frac{k}{2}$个元素。随后我们可以继续用这个方法排除下去，直到k=1。

当然，有一些特殊情况需要考虑：

• a[k/2 - 1]和b[k/2 - 1]越界，我们需要选取的是数组的最后一个元素。
• 其中一个数组的元素全部被排除，可以直接返回另一个数组第$k$小的元素。
• $k=1$，只要返回第一个元素的较小值即可。

// C++
// https://leetcode.cn/problems/median-of-two-sorted-arrays/solutions/258842/xun-zhao-liang-ge-you-xu-shu-zu-de-zhong-wei-s-114/
class Solution {
public:
    int getKthElement(const vector<int>& nums1, const vector<int>& nums2, int k) {
        /* 主要思路：要找到第 k (k>1) 小的元素，那么就取 pivot1 = nums1[k/2-1] 和 pivot2 = nums2[k/2-1] 进行比较
         * 这里的 "/" 表示整除
         * nums1 中小于等于 pivot1 的元素有 nums1[0 .. k/2-2] 共计 k/2-1 个
         * nums2 中小于等于 pivot2 的元素有 nums2[0 .. k/2-2] 共计 k/2-1 个
         * 取 pivot = min(pivot1, pivot2)，两个数组中小于等于 pivot 的元素共计不会超过 (k/2-1) + (k/2-1) <= k-2 个
         * 这样 pivot 本身最大也只能是第 k-1 小的元素
         * 如果 pivot = pivot1，那么 nums1[0 .. k/2-1] 都不可能是第 k 小的元素。把这些元素全部 "删除"，剩下的作为新的 nums1 数组
         * 如果 pivot = pivot2，那么 nums2[0 .. k/2-1] 都不可能是第 k 小的元素。把这些元素全部 "删除"，剩下的作为新的 nums2 数组
         * 由于我们 "删除" 了一些元素（这些元素都比第 k 小的元素要小），因此需要修改 k 的值，减去删除的数的个数
         */
 
        int m = nums1.size();
        int n = nums2.size();
        int index1 = 0, index2 = 0;
 
        while (true) {
            // 边界情况
            if (index1 == m) {
                return nums2[index2 + k - 1];
            }
            if (index2 == n) {
                return nums1[index1 + k - 1];
            }
            if (k == 1) {
                return min(nums1[index1], nums2[index2]);
            }
 
            // 正常情况
            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];
            if (pivot1 <= pivot2) {
                k -= newIndex1 - index1 + 1;
                index1 = newIndex1 + 1;
            }
            else {
                k -= newIndex2 - index2 + 1;
                index2 = newIndex2 + 1;
            }
        }
    }
 
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int totalLength = nums1.size() + nums2.size();
        if (totalLength % 2 == 1) {
            return getKthElement(nums1, nums2, (totalLength + 1) / 2);
        }
        else {
            return (getKthElement(nums1, nums2, totalLength / 2) + getKthElement(nums1, nums2, totalLength / 2 + 1)) / 2.0;
        }
    }
};