寻找两个正序数组的中位数

题目描述

给定两个大小分别为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。请你找出并返回这两个正序数组的 中位数。

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

示例

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

我的解题过程

第一次尝试 - 直接合并排序

第一次想的比较简单，就是把两个数组合并起来排个序，然后判断一下总长度是奇数还是偶数，再拿到对应的下标去计算中位数。

function findMedianSortedArrays(nums1: number[], nums2: number[]): number {
  const arr = nums1.concat(nums2).sort((a, b) => a - b);
  const n = arr.length;
  const isOdd = n % 2 === 1;
  const index = Math.floor(n / 2);

  if (isOdd) {
    return arr[index];
  }

  return (arr[index - 1] + arr[index]) / 2;
}

这个方法能过，但时间复杂度是 O((m+n)log(m+n))，不符合题目要求的 O(log(m+n))。

学习二分法 - 先做 704 题打基础

看到题目要求 O(log(m+n))，想到应该是要用二分法。但我对二分法还不太熟，就去刷了 704 题（二分查找）先练练手。

二分法主要有两种写法，区别在于区间的定义：

第一种：闭区间 [left, right]

function search(nums: number[], target: number): number {
  let left = 0;
  let right = nums.length - 1;

  while (left <= right) {
    const mid = Math.floor((left + right) / 2);

    if (nums[mid] > target) {
      right = mid - 1;
    } else if (nums[mid] < target) {
      left = mid + 1;
    } else {
      return mid;
    }
  }
  return -1;
}

第二种：左闭右开 [left, right)

function search(nums: number[], target: number): number {
  let left = 0;
  let right = nums.length;

  while (left < right) {
    const mid = Math.floor((left + right) / 2);

    if (nums[mid] > target) {
      right = mid;
    } else if (nums[mid] < target) {
      left = mid + 1;
    } else {
      return mid;
    }
  }
  return -1;
}

理解了二分法的基本思路后，就可以回来做这道题了。

最终解法 - 在较短数组上二分

这道题的核心思路是：不用真的合并数组，而是通过二分法找到一个分割点，把两个数组分成”左半部分”和”右半部分”，保证：

1. 左半部分的元素个数等于右半部分（或差 1）
2. 左半部分的最大值 ≤ 右半部分的最小值

这样中位数就是左半部分的最大值和右半部分的最小值算出来的。

function findMedianSortedArrays(nums1, nums2) {
  if (nums1.length > nums2.length) {
    return findMedianSortedArrays(nums2, nums1);
  }

  const m = nums1.length;
  const n = nums2.length;

  let left = 0;
  let right = m;

  const halfLen = Math.floor((m + n + 1) / 2);

  while (left <= right) {
    const i = Math.floor((left + right) / 2);
    const j = halfLen - i;

    const nums1LeftMax = i === 0 ? -Infinity : nums1[i - 1];
    const nums1RightMin = i === m ? Infinity : nums1[i];

    const nums2LeftMax = j === 0 ? -Infinity : nums2[j - 1];
    const nums2RightMin = j === n ? Infinity : nums2[j];

    if (nums1LeftMax <= nums2RightMin && nums2LeftMax <= nums1RightMin) {
      if ((m + n) % 2 === 1) {
        return Math.max(nums1LeftMax, nums2LeftMax);
      }

      return (Math.max(nums1LeftMax, nums2LeftMax) + Math.min(nums1RightMin, nums2RightMin)) / 2;
    } else if (nums1LeftMax > nums2RightMin) {
      right = i - 1;
    } else {
      left = i + 1;
    }
  }

  // 理论上不会走到这里
  throw new Error("Input arrays are not sorted.");
}

小结

这道题是 Hard 难度，确实不简单。几个关键点：

1. 为什么要在短数组上二分：为了保证另一个数组的分割点 j 不会越界（j = halfLen - i 必须 ≥ 0）
2. 如何确定分割是否正确：左半部分的最大值要 ≤ 右半部分的最小值，即 nums1[i-1] <= nums2[j] 且 nums2[j-1] <= nums1[i]
3. 边界处理：用 -Infinity 和 Infinity 处理数组边界，避免复杂的条件判断
4. 奇偶处理：总长度为奇数时取左半部分最大值，偶数时取左右部分临界值的平均

这道题我是先做了 704 题打基础，理解了二分法的思路后才做出来的。如果直接上来做这道 Hard 题可能会比较懵，建议先把基础的二分查找掌握好。