寻找两个正序数组的中位数
这是一道比较难的问题,LeetCode4.给定两个大小分别为 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
对于本题,最直观的思路有以下两种:
- 使用归并的方式,合并两个有序数组,得到一个大的有序数组。大的有序数组的中间位置的元素,即为中位数。这种方式的时间复杂度是O(m+n),空间复杂度是 O(m+n)。
- 另一种方式,不需要合并两个有序数组,只要找到中位数的位置即可。由于两个数组的长度已知,因此中位数对应的两个数组的下标之和也是已知的。维护两个指针,初始时分别指向两个数组的下标 00 的位置,每次将指向较小值的指针后移一位(如果一个指针已经到达数组末尾,则只需要移动另一个数组的指针),直到到达中位数的位置。这种方式可以将空间复杂度降到 O(1),但是时间复杂度仍是 O(m+n)。
如何把时间复杂度降低到O(log(m+n)) 呢?如果对时间复杂度的要求有log,通常都考虑二分、快排或者堆三个方面。而对于有序序列,通常要考虑一下能否使用二分来解决。
如果要使用二分,核心问题是基于什么规则将数据砍掉一半。而本题是两个序列,所以我们的核心问题是如何从两个序列中分别砍半,图示如下k=(m+n)/2:
根据中位数的定义,当m+n
是奇数时,中位数是两个有序数组中的第(m+n)/2
个元素,当m+n
是偶数时,中位数是两个有序数组中的第(m+n)/2
个元素和第(m+n)/2+1
个元素的平均值。因此,这道题可以转化成寻找两个有序数组中的第 k
小的数,其中 k
为 (m+n)/2
或 (m+n)/2+1
。
假设两个有序数组分别是LA
和LB
。要找到第 k
个元素,我们可以比较LA[k/2−1]
和LB[k/2−1]
。由于LA[k/2−1]
和LB[k/2−1]
的前面分别有LA[0..k/2−2]
和 LB[0..k/2−2]
,即k/2−1
个元素,对于LA[k/2−1]
和LB[k/2−1]
中的较小值,最多只会有(k/2−1)+(k/2−1)≤k−2
个元素比它小,那么它就不能是第 k
小的数了。
因此我们可以归纳出以下几种情况:
-
如果
LA[k/2−1] < LB[k/2−1]
,则比LA[k/2−1]
小的数最多只有LA
的前k/2−1
个数和LB
的前k/2−1
个数,即比LA[k/2−1]
小的数最多只有k−2
个,因此LA[k/2−1]
不可能是第k
个数,LA[0]
到LA[k/2−1]
也都不可能是第k
个数,可以全部排除。 -
如果
LA[k/2−1] > LB[k/2−1]
,则可以排除LB[0]
到LB[k/2−1]
。也就是一次砍掉一半。 -
如果
LA[k/2−1] = LB[k/2−1]
,则可以归入第一种情况处理。
可以看到,比较LA[k/2−1]
和LB[k/2−1]
之后,可以排除k/2
个不可能是第 k
小的数,查找范围缩小了一半。同时,我们将在排除后的新数组上继续进行二分查找,并且根据我们排除数的个数,减少 k 的值,这是因为我们排除的数都不大于第 k 小的数。
以下边界情况需要特殊处理以下:
-
如果
LA[k/2−1]
或者LB[k/2−1]
越界,那么我们可以选取对应数组中的最后一个元素。在这种情况下,我们必须根据排除数的个数减少k
的值,而不能直接将k
减去k/2
。 -
如果
k=1
,我们只要返回两个数组首元素的最小值即可。
实现代码:
class Solution {
public double findMedianSortedArrays(int[] nums1, int[] nums2) {
int length1 = nums1.length, length2 = nums2.length;
int totalLength = length1 + length2;
if (totalLength % 2 == 1) {
int midIndex = totalLength / 2;
double median = getKthElement(nums1, nums2, midIndex + 1);
return median;
} else {
int midIndex1 = totalLength / 2 - 1, midIndex2 = totalLength / 2;
double median = (getKthElement(nums1, nums2, midIndex1 + 1) + getKthElement(nums1, nums2, midIndex2 + 1)) / 2.0;
return median;
}
}
public int getKthElement(int[] nums1, int[] nums2, int k) {
int length1 = nums1.length, length2 = nums2.length;
int index1 = 0, index2 = 0;
int kthElement = 0;
while (true) {
// 边界情况
if (index1 == length1) {
return nums2[index2 + k - 1];
}
if (index2 == length2) {
return nums1[index1 + k - 1];
}
if (k == 1) {
return Math.min(nums1[index1], nums2[index2]);
}
// 正常情况
int half = k / 2;
int newIndex1 = Math.min(index1 + half, length1) - 1;
int newIndex2 = Math.min(index2 + half, length2) - 1;
int pivot1 = nums1[newIndex1], pivot2 = nums2[newIndex2];
if (pivot1 <= pivot2) {
k -= (newIndex1 - index1 + 1);
index1 = newIndex1 + 1;
} else {
k -= (newIndex2 - index2 + 1);
index2 = newIndex2 + 1;
}
}
}
}