寻找两个有序数组的中位数
非常好的讲解
题目
给定两个大小为 m 和 n 的有序数组 nums1 和 nums2。
请你找出这两个有序数组的中位数,并且要求算法的时间复杂度为 O(log(m + n))。
你可以假设 nums1 和 nums2 不会同时为空。
示例1:
nums1 = [1, 3]
nums2 = [2]
则中位数是 2.0
示例2:
nums1 = [1, 2]
nums2 = [3, 4]
则中位数是 (2 + 3)/2 = 2.5
解题思路分析
这道题的直观解法:就是扫一遍两个数组,找到中间位置的那一个/两个数,然后得到中位数。时间复杂度是 O(m+n),其中 m 和 n 分别是数组 A 和数组 B 的长度。
但是!面试官会这么轻易就放过你吗?显然是不可能滴~~我偷看了一下题目描述下的“challenge”标签,原来这道题的最优解是 O(log(m + n)) 的复杂度。(m + n) 是俩数组合并后的总长度 L,看到 O(log L) 这样的复杂度,而且还是有序数组,能想到哪个算法吗?没错,就是二分查找!那我们试试。。
如果我取出 A[i] ,用二分查找在数组 B 中找到 A[i] 可以插入的位置,假设 A[i] 在 B 中的插入位置是 j,那么 A[i] 在整个合并数组中的位置就是 (i + j) ,因为要求的中位数的位置是 (m + n) / 2,通过比较 (i + j) 和 (m + n) / 2 的大小可以每次舍弃 A 的一部分,从而收敛数组 A。用同样的方法可以收敛数组 B。但是这样的复杂度是 O(log m * log n),复杂度大于 O(log(m + n)),显然不是最优的。
那如果不是直接使用二分查找算法,而是借用二分查找的思想呢?就是每次有选择地舍弃掉数组的一部分,从而达到收敛数组缩小查找范围的效果。
a. 我们重新看题目,要找中位数,就是要找第 k 大的数(k = (L / 2 + 1),其中L 是上面提到的合并后新数组的长度,当 L 是偶数时,要求第 (L / 2) 大和第 (L / 2 + 1) 大的两个数)。当我们舍弃掉一部分,假设舍弃部分的长度为 length,那么接下来就是在剩下的数组里求第 (k - length) 大的数。逐层缩小范围,直到两数组其中一个走完,或者要求的是第 1 大的元素,就可以直接返回结果了。
b. 那如何“选择”要舍弃哪部分呢?既然是要找合并后的数组 C 的第 k 大元素,即 C[k-1],那如果我们从 A 和 B 中分别取前 k/2 个元素,其中必然有一部分是是在数组 C 的前 k 个数里。设 mid = k / 2,当 A[mid - 1] < B[mid - 1] 时,可以断定 A 的前 mid 个元素是在 C 的前 k 个数里(此处可用反证法得证),那么我们则舍弃 A 的前 mid 个元素。反之则舍弃 B 的前 mid 个元素。现在数组 A 或者 B 已经舍弃掉 k/2 个元素,缩小查找范围了,那接下来可以按照同样的方法继续选择吗?当然!现在剩下总共 (L - mid) 个元素,且 A 和 B 依旧有序,要找的是第 (k - mid) 大的元素,所以我们可以按照上面的方法继续递归选择下去,直到找到目标元素!
c. 复杂度分析:每次从合并后数组 C 里减少 k/2 个元素,直到找到目标元素。所以时间复杂度是 O(log L) = O(log (m + n)) !
class Solution {
public double findMedianSortedArrays(int[] nums1, int[] nums2) {
/**
如果两个数组的中位数 mid1 < mid2, 则说明合并后的中位数位于 num1.right + num2之间
否则合并后的中位数位于 nums2.right + nums1 之间 (right 是相对于 mid 而言的)
getKth 函数负责找到两个数组合并(假设)后有序的数组中的第 k 个元素, k 从 1 开始计算
**/
if(nums1.length == 0 && nums2.length == 0) return 0.0;
int m = nums1.length, n = nums2.length;
// l: 合并后数组的左半部分的最后一个数 r: 合并后数组的右半部分的第一个数
int l = (m+n+1) / 2;
int r = (m+n+2) / 2;
// 如果 m+n 是奇数 getKth 的返回值是相同的, 是偶数则是合并后数组的中间两个数
if(l == r) return getKth(nums1, 0, nums2, 0, l);
return (getKth(nums1, 0, nums2, 0, l) + getKth(nums1, 0, nums2, 0, r)) / 2.0;
}
private double getKth(int[] nums1, int st1, int[] nums2, int st2, int k) {
// 边界情况, 如果 nums1数组已经穷尽了, 则只能返回 nums2 中的第 k 个元素
if(st1 > nums1.length-1) return nums2[st2 + k - 1];
if(st2 > nums2.length-1) return nums1[st1 + k - 1];
// 边界情况, k = 1 则返回两个数组中最小的那个
if(k == 1) return Math.min(nums1[st1], nums2[st2]);
// 在 nums1 和 nums2 当前范围内找出 mid1 和 mid2 判断舍弃哪半部分
int mid1 = Integer.MAX_VALUE;
int mid2 = Integer.MAX_VALUE;
if(st1 + k/2 - 1 < nums1.length) mid1 = nums1[st1 + k/2 - 1];
if(st2 + k/2 - 1 < nums2.length) mid2 = nums2[st2 + k/2 - 1];
// mid1 < mid2 在 nums1.right 和 nums2 之间搜索, 丢掉 k/2 个数.
if(mid1 < mid2)
return getKth(nums1, st1 + k/2, nums2, st2, k - k/2);
else
return getKth(nums1, st1, nums2, st2 + k/2, k - k/2);
}
}