算法讲解 —— 寻找两个升序数组的中位数

前言

明天考试，但是实在没有复习的欲望，感觉丝毫没有前几天刷 leetcode 有意思，姑且就写一篇算法讲解的 blog，讲一下我觉得非常有意思的一道题：

4. 寻找两个升序数组的中位数
给定两个大小分别为和的正序（从小到大）数组和。请你找出并返回这两个正序数组的中位数，算法的时间复杂度应该为。

分析

看到有序 + 复杂度，显然这是一个二分查找的题，但是难点就在于这里同时有两个数组，如果做归并（merge）并提前返回，时间复杂度一定是的，所以我们必须在不合并的前提下同时对两个数组做二分查找。

而且，令总数组长度为，对于长度不同的数组，我们所要二分查找的数还不一样：

对于奇数长度的数组，我们仅需要查找索引处的数即可。
对于偶数长度的数组，我们需要查找和索引处的数相加除以。

思路

首先，我们可以抽象出一个过程（Procedure），给定两个有序数组，找到索引为的元素，函数原型类似于：

int findKth(const std::vector<int>& nums1, const std::vector<int>& nums2, int k);

如果我们实现好了这个函数，那么我们最终的答案很简单的就是：

class Solution {
public:
    double findMedianSortedArrays(std::vector<int>& nums1, std::vector<int>& nums2) {
        int m = nums1.size(), n = nums2.size();
        if ((m + n) % 2 == 0) {
            return (double) (
              findKth(nums1, nums2, (m + n) / 2 - 1) +
              findKth(nums1, nums2, (m + n) / 2)
            ) / 2;
        } else {
            return findKth(nums1, nums2, (m + n) / 2);
        }
    }
};

因此接下来我们需要实现 findKth 这个函数。

然后，令这两个有序数组分别为和，归并后的数组为，并且不失一般性的，不妨设，以便后续分析。

对于，我们要找的数肯定要么位于，要么位于，要么作为重复元素同时位于和。不管怎样，我们都可以假设，在通过二分查找到索引为的数时，在中至少有个数小于等于这个数，在中至少有个数小于等于这个数，则有。

证明很简单，如果我们在合并后的数组中根据索引将其切成两半，那么左边正好有个数。由于和各自有序，这个切分对应到两个数组中也一定是两个前缀，设这两个前缀长度分别为和，自然有。

举个例子，，我找的是的元素，归并后找一下发现是这个数。在左侧，贡献了这个数，贡献了这个数。因此有。

我们二分的时候所需要找的就是合适的和，正好对应着两个数组中位于左侧的前缀长度，而或者所在索引处正好对应着我们需要的索引为的元素，即。

因此，一个初步的二分模板如下所示：

int findKth(const std::vector<int>& nums1, const std::vector<int>& nums2, int k) {
  int m = nums1.size(), n = nums2.size();
  int l = std::max(0, k - n), r = std::min(m, k);
  while (l <= r) {
    int k1 = l + (r - l) / 2;
    int k2 = k - k1;

    // TODO
  }
}

其中 k2 由 k1 被 k 减去而得到，因此天然保证了的约束，同时为了防止 k 越界，并且尽可能预先缩小二分区间，初始化 l 和 r 的时候相应的做了调整。

但是，我们怎么知道什么情况下和恰好是左侧的前缀长度呢？这时候，我们就要思考一下如果恰好是的情况下，会有什么特性：

两数组有序，因此显然且。
，因此显然且。
- ，因此可证，否则，如果，则有，那么在中，比小的数至少有个（从到），而在中，比小的数有个（从到），因此在中，比小的数至少有个，但是，矛盾！
- ，同理可证。

纯数学语言可能不是很直观，画出来就很直观了。

如果，那么这种错误情况下的错误排列如下所示：

而和恰好是左侧的前缀长度时，对应的和数组中元素的正确排列方式为：

其中，或者

其中。

因此，在二分查找过程中，只要且，那么我们就找到了正确的和，这便是二分查找的退出条件。

而如果，那么说明太小了，而太大了；如果，那么说明太小了，而太大了，这便是二分查找缩小边界的规则。

因此，我们可以有如下实现：

int findKth(const std::vector<int>& nums1, const std::vector<int>& nums2, int k) {
  int m = nums1.size(), n = nums2.size();
  int l = std::max(0, k - n), r = std::min(m, k);
  while (l <= r) {
    int k1 = l + (r - l) / 2;
    int k2 = k - k1;

    int l1 = (k1 == 0) ? INT_MIN : nums1[k1 - 1];
    int r1 = (k1 == m) ? INT_MAX : nums1[k1];
    int l2 = (k2 == 0) ? INT_MIN : nums2[k2 - 1];
    int r2 = (k2 == n) ? INT_MAX : nums2[k2];

    if (l1 <= r2 && l2 <= r1) {
      return std::min(r1, r2);
    } else if (l1 > r2) {
      r = k1 - 1;
    } else {
      l = k1 + 1;
    }
  }

  return -1; // not found
}

工程上，因为可能会有数组越界的问题，所以如果遇到数组边界，我们就将边界设置为 INT_MIN 和 INT_MAX。

综上，这题的一种解法如下所示：

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

    int m = nums1.size(), n = nums2.size();
    if ((m + n) % 2 == 0) {
      return (double) (
        findKth(nums1, nums2, (m + n) / 2 - 1) +
        findKth(nums1, nums2, (m + n) / 2)
      ) / 2;
    } else {
      return findKth(nums1, nums2, (m + n) / 2);
    }
  }
private:
  int findKth(const std::vector<int>& nums1, const std::vector<int>& nums2, int k) {
    int m = nums1.size(), n = nums2.size();
    int l = std::max(0, k - n), r = std::min(m, k);
    while (l <= r) {
      int k1 = l + (r - l) / 2;
      int k2 = k - k1;

      int l1 = (k1 == 0) ? INT_MIN : nums1[k1 - 1];
      int r1 = (k1 == m) ? INT_MAX : nums1[k1];
      int l2 = (k2 == 0) ? INT_MIN : nums2[k2 - 1];
      int r2 = (k2 == n) ? INT_MAX : nums2[k2];

      if (l1 <= r2 && l2 <= r1) {
        return std::min(r1, r2);
      } else if (l1 > r2) {
        r = k1 - 1;
      } else {
        l = k1 + 1;
      }
    }
    return -1;
  }
};

其中，为了加快二分收敛，所以我们通过交换引用从而尽可能减少被二分的 nums1 的长度。

优化

其实，在这种解法之上，我们还可以再做一些优化，让 findKth 只需要调用一次！

注意到，我们在使用二分查找的时候，额外选用的是和这两个左边的索引，那么显然，如果我用和这两个右边的索引，然后把函数里面对应的符号修改一下，是不是也是一样的能够实现 findKth 函数？

再进一步，实际上 findKth 函数不仅找到了索引为的数，同时还找到了索引为的数！即，，且！

证明也不难，把之前列举的那几种和的相对位置的情况全都重写一遍就结束了。

因此，优化后的解法如下所示：

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

    int m = nums1.size(), n = nums2.size();
    int l = 0, r = m;

    while (l <= r) {
      int k1 = l + (r - l) / 2;
      int k2 = (m + n + 1) / 2 - k1;

      int l1 = (k1 == 0) ? INT_MIN : nums1[k1 - 1];
      int r1 = (k1 == m) ? INT_MAX : nums1[k1];
      int l2 = (k2 == 0) ? INT_MIN : nums2[k2 - 1];
      int r2 = (k2 == n) ? INT_MAX : nums2[k2];

      if (l1 <= r2 && l2 <= r1) {
        if ((m + n) % 2 == 0) {
          return (double) (std::max(l1, l2) + std::min(r1, r2)) / 2;
        } else {
          return std::max(l1, l2);
        }
      } else if (l1 > r2) {
        r = k1 - 1;
      } else {
        l = k1 + 1;
      }
    }

    return -1; // not found
  }
};

其中，我们令，对于偶数长度的数组来说，，而，正好是求中位数所需的两个数；对于奇数长度的数组来说，，而，后者是我们求中位数所需的那个数。因此，不论是偶数还是奇数长度，我们都可以在一次二分查找中全部找到所需的数。

结语

这是一道非常经典并且有趣的二分查找的困难题，很久之前就在我校《算法基础》课程的考试中出现过，当时我实在没想出来有什么好的解法，直到最近刷 leetcode 才想到如此优雅高效并且易于理解的好方法，希望能帮助读者更好吃透这道优秀的题目。

前言

分析

思路#

优化#

结语#

思路

优化

结语