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

题目描述

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

解题思路

有序数组中找中位数，其实就是找特定下标的元素。

简单的思路是先将两个有序数组归并，归并后总元素数量为 m+n ，

  • 若总元素数目为奇数，则寻找下标为 (m+n)/2 的元素
  • 若元素为偶数，则寻找下标为 (m+n)/2-1 和 (m+n)/2 的元素，取平均

该思路的时间复杂度为 O(m+n) 。

如果省去归并操作，直接使用二分法来寻找特定的元素，那么就可以把事件复杂度降至 O(log(m+n)) 。

这种思路的关键在于设计二分法的策略。首先，“两个正序数组中寻找中位数”可以转化为“两个正序数组中寻找第k大的数”，然后，使用排除法，一点一点地去掉第1~k-1大的数。

二分法的循环体

比较 nums1[k/2 - 1] 和 nums2[k/2 - 1] ，如果 nums1[k/2 - 1] <= nums2[k/2 - 1] ：

  • nums2[k/2 - 1] >= nums1[k/2 - 1] >= ... >= nums1[0] ，排除掉 k/2 个元素

在下一轮循环中，nums1[0] 、nums1[1] 、…、nums1[k/2 - 1] 被排除，并且 k -= k/2 。

二分法的边界条件

  • 如果某数组为空，那么返回另一个数组第k大的数
  • 如果 k=1 ，那么返回两数组首元素中较小的数

二分法的特殊情况

每轮循环中，我们设置的下标 index + k/2 - 1 可能越界，此时需要把二分法更新迭代的步长设为0。

复杂度分析

时间复杂度：O(log(m+n))
空间复杂度：O(1)

知识点

二分法、中位数

代码

  1. class Solution {
  2. public:
  3.     double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
  4.         int size1 = nums1.size();
  5.         int size2 = nums2.size();
  6.         if ((size1 + size2) % 2 == 1) {
  7.             return findKthElement(nums1, nums2, (size1 + size2)/2 + 1);
  8.         } else {
  9.             cout << findKthElement(nums1, nums2, (size1 + size2)/2) << endl;
  10.             cout << findKthElement(nums1, nums2, (size1 + size2)/2 + 1) << endl;
  11.             return (findKthElement(nums1, nums2, (size1 + size2)/2)
  12.                 + findKthElement(nums1, nums2, (size1 + size2)/2 + 1)) / 2.0;
  13.         }
  14.     }
  16.     double findKthElement(vector<int>& nums1, vector<int>& nums2, int k) {
  17.         int begIndex1 = 0;
  18.         int begIndex2 = 0;
  19.         int size1 = nums1.size();
  20.         int size2 = nums2.size();
  21.         int step1 = k/2 - 1;
  22.         int step2 = k/2 - 1;
  24.         while (true) {
  25.             if (size1 - begIndex1 == 0) {
  26.                 return nums2[begIndex2 + k - 1];
  27.             } else if (size2 - begIndex2 == 0) {
  28.                 return nums1[begIndex1 + k - 1];
  29.             } else if (k == 1) {
  30.                 return min(nums1[begIndex1], nums2[begIndex2]);
  31.             } else if (begIndex1 + step1 >= size1) {
  32.                 step1 = 0;
  33.             } else if (begIndex2 + step2 >= size2) {
  34.                 step2 = 0;
  35.             }
  37.             if (nums1[begIndex1 + step1] <= nums2[begIndex2 + step2]) {
  38.                 begIndex1 += step1 + 1;
  39.                 // begIndex2 += step2;
  40.                 k -= step1 + 1;
  41.             } else {
  42.                 // begIndex1 += step1;
  43.                 begIndex2 += step2 + 1;
  44.                 k -= step2 + 1;
  45.             }
  46.             step1 = k/2 - 1;
  47.             step2 = k/2 - 1;
  49.             // cout << "begIndex1: " << begIndex1 << endl;
  50.             // cout << "begIndex2: " << begIndex2 << endl;
  51.             // cout << "k: " << k << endl;
  52.         }
  53.     }
  54. };

参考资料

  1. Leetcode官方题解