[二分搜索|快速选择] leetcode 4 寻找两个正序数组的中位数
1.题目
题目链接
给定两个大小为 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
示例3:
输入:nums1 = [0,0], nums2 = [0,0]
输出:0.00000
示例4:
输入:nums1 = [], nums2 = [1]
输出:1.00000
示例5:
输入:nums1 = [2], nums2 = []
输出:2.00000
2.分析
2.1.直观思路
由于给出的两个数组是有序的,因此可以直接将两个数组合并形成一个新的有序数组。
这一过程与归并排序中“并”的过程十分相似:
double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
vector<int> nums3; //合并后的新数组
int len1 = nums1.size();
int len2 = nums2.size();
int i = 0, j = 0;
while(i < len1 || j < len2) { //遍历原两数组
//如果数组1已经遍历完毕,将数组2的内容拷贝至新数组
if(i > len1 - 1) {
while(j < len2){
nums3.push_back(nums2[j]);
j++;
}
break;
}
//如果数组2已经遍历完毕,将数组1的内容拷贝至新数组
if(j > len2 - 1) {
while(i < len1){
nums3.push_back(nums1[i]);
i++;
}
break;
}
//选取数组1,2中较小的那个元素加入新数组
if(nums1[i] <= nums2[j]) {
nums3.push_back(nums1[i]);
i++;
}else{
nums3.push_back(nums2[j]);
j++;
}
}
//取中位数
if(nums3.size() % 2 != 0) {
return nums3[nums3.size() / 2] * 1.0;
}else{
return (nums3[nums3.size() / 2] + nums3[nums3.size() / 2 - 1]) / 2.0;
}
}
此时时间复杂度与空间复杂度均为O(m+n)。
实际上,由于两个原数组的长度是已知的,因此合并后新数组中位数的位置也是已知的,并不需要新数组num3,只需要一个变量index统计当前已经选择多少元素即可。这样可以将空间复杂度降为O(1)。
2.2.改进
题目要求的复杂度为O(log(m + n)),为对数级,似乎可以尝试使用二分的思想来解决。
为此,可以将“选取中位数”问题视为一个经典的选取数组中第K小元素的问题,即:
double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
int len = nums1.size() + nums2.size();
if (len % 2 != 0) {
return getKthElement(nums1, nums2, (len + 1) / 2);
} else {
return (getKthElement(nums1, nums2, len / 2) + getKthElement(nums1, nums2, len / 2 + 1)) / 2.0;
}
}
2.3.如何选取第K小元素
思路:令pivot1 = nums1[k / 2 - 1],pivot2 = nums2[k / 2 - 1]。取pivot = min(pivot1, pivot2)。因此num1与nums2中比pivot小的元素最多有k / 2 - 1 + k / 2 - 1 = k - 2个,这样就可以确保pivot最多为第k - 1小的元素(甚至更小)。
如果pivot = pivot1,那么pivot及之前的数都不可能为第k小的元素(因为他们更小),因此我们可以排除nums1[0~k/2-1]这些元素。
想象他们在数组中被删除,那么目前需要寻找的就是第“k - k / 2”小的元素了。这样就减小的问题的规模,直到最后“k = 1”,此时第一个元素便是我们要找的第K小元素。
3.代码
class Solution {
public:
int getKthElement(const vector<int>& nums1, const vector<int>& nums2, int k) {
int m = nums1.size();
int n = nums2.size();
//index之前的元素都是我们“已经排除掉的”
int index1 = 0, index2 = 0;
while (true) {
//如果nums1中的元素全部排除了,那么我们选取nums2“排除后”的第k个元素即可
if (index1 == m) {
return nums2[index2 + k - 1];
}
//如果nums2中的元素全部排除了,那么我们选取nums1“排除后”的第k个元素即可
if (index2 == n) {
return nums1[index1 + k - 1];
}
//k=1则直接选取两数组中最小的那个元素返回即可
if (k == 1) {
return min(nums1[index1], nums2[index2]);
}
//选取pivot元素(注意越界)
int newIndex1 = min(index1 + k / 2 - 1, m - 1);
int newIndex2 = min(index2 + k / 2 - 1, n - 1);
int pivot1 = nums1[newIndex1];
int pivot2 = nums2[newIndex2];
//nums1中的元素更小,那么k需要减去nums1中被排除的第index1个元素至第newIndex1个元素
if (pivot1 <= pivot2) {
k -= newIndex1 - index1 + 1;
index1 = newIndex1 + 1; //index1之前的元素都被排除了
}else { //反之同理
k -= newIndex2 - index2 + 1;
index2 = newIndex2 + 1;
}
}
}
double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
int len = nums1.size() + nums2.size();
if (len % 2 == 1) {
return getKthElement(nums1, nums2, (len + 1) / 2);
} else {
return (getKthElement(nums1, nums2, len / 2) + getKthElement(nums1, nums2, len / 2 + 1)) / 2.0;
}
}
};