题目
给定两个大小为 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
来源:力扣(LeetCode)
思路
首先,假设两个有序数组的长度分别为m和n,由于两个数组长度之和 m+n 的奇偶性不确定,因此需要分情况来讨论:对于奇数的情况,直接找到最中间的数即可,偶数的话需要求最中间两个数的平均值。为了简化代码,不分情况讨论,使用一个技巧,分别找第 (m+n+1) / 2
个,和 (m+n+2) / 2
个,然后求其平均值,这对奇偶数都是适用的。若 m+n
为奇数的话,那么 (m+n+1) / 2
和 (m+n+2) / 2
的值相等,相当于两个相同的数字相加再除以2,还是其本身。
其次,定义一个函数在两个有序数组中找到第K个元素,如何实现找到第K个元素?首先,为了避免拷贝产生新的数组从而增加时间复杂度,使用两个变量i
和j
分别来标记数组 nums1
和 nums2
的起始位置。然后来处理一些小问题,比如当某一个数组的起始位置大于等于其数组长度时,说明其所有数字均已经被淘汰,相当于一个空数组,实际上就变成了在另一个数组中找数字,当然直接就可以找出来。还有就是,如果 K=1
的话,只要比较 nums1
和 nums2
的起始位置i
和j
上的数字就可以了。难点就在于一般的情况怎么处理?因为需要在两个有序数组中找到第K个元素,为了加快搜索的速度,可以使用二分法,本题就是要对K
二分,意思是需要分别在 nums1
和 nums2
中查找第 K/2
个元素,注意,这里由于两个数组的长度不定,所以有可能某个数组没有第 K/2
个数字,所以需要先 检查 一下,数组中到底存不存在第 K/2
个数字,如果存在就取出来,否则就赋值一个整型最大值(目的是要在 nums1
或者 nums2
中先淘汰 K/2
个较小的数字,判断的依据就是看 midVal1
和 midVal2
谁更小,但如果某个数组的个数不到 K/2
个,自然无法淘汰,所以将其对应的 midVal
值设为整型最大值,以保证其不会被淘汰),若某个数组没有第 K/2
个数字,则淘汰另一个数组的前 K/2
个数字即可。
举例说明,比如 nums1
= {3},nums2
= {2, 4, 5, 6, 7},K=4,要找两个数组混合中第4个数字,则分别在 nums1
和 nums2
中找第2个数字,而 nums1
中只有一个数字,不存在第二个数字,则 nums2
中的前2个数字可以直接跳过,为啥呢,因为要求的是整个混合数组的第4个数字,不管 nums1
中的那个数字是大是小,第4个数字绝不会出现在 nums2
的前两个数字中,所以可以直接跳过。
最后,有没有可能两个数组都不存在第 K/2
个数字呢,这道题里是不可能的,因为K
不是任意给的,而是给的 m+n
的中间值,所以必定至少会有一个数组是存在第 K/2
个数字的。最后就是二分法的核心,比较这两个数组的第 K/2
小的数字 midVal1
和 midVal2
的大小,如果第一个数组的第 K/2
个数字小的话,那么说明要找的数字肯定不在 nums1
中的前 K/2
个数字,可以将其淘汰,将 nums1
的起始位置向后移动 K/2
个,并且此时的K也自减去 K/2
,调用递归,举个例子来说,比如 nums1
= {1, 3},nums2
= {2, 4, 5},K=4
,要找两个数组混合中第4个数字,那么分别在 nums1
和 nums2
中找第2个数字,nums1
中的第2个数字是3,nums2
中的第2个数字是4,由于3小于4,所以混合数组中第4个数字肯定在 nums2
中,可以将 nums1
的起始位置向后移动 K/2 个。反之,淘汰 nums2
中的前 K/2
个数字,并将 nums2
的起始位置向后移动 K/2
个,并且此时的K也自减去 K/2
,调用递归即可。
C++代码
class Solution {
public:
double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2)
{
int m = nums1.size(), n = nums2.size(), left = (m + n + 1) / 2, right = (m + n + 2) / 2;
return (findKth(nums1, 0, nums2, 0, left) + findKth(nums1, 0, nums2, 0, right)) / 2.0;
}
int findKth(vector<int>& nums1, int i, vector<int>& nums2, int j, int k)
{
if (i >= nums1.size())
return nums2[j + k - 1];
if (j >= nums2.size())
return nums1[i + k - 1];
if (k == 1)
return min(nums1[i], nums2[j]);
int midVal1 = (i + k / 2 - 1 < nums1.size()) ? nums1[i + k / 2 - 1] : INT_MAX;
int midVal2 = (j + k / 2 - 1 < nums2.size()) ? nums2[j + k / 2 - 1] : INT_MAX;
if (midVal1 < midVal2)
{
return findKth(nums1, i + k / 2, nums2, j, k - k / 2);
}
else
{
return findKth(nums1, i, nums2, j + k / 2, k - k / 2);
}
}
};