原题为:
There are two sorted arrays nums1 and nums2 of size m and n respectively. Find the median of the two sorted arrays. The overall run time complexity should be O(log (m+n)).
好吧,我是普通青年,第一反应就是使用Merge两个数字的思想,直接找到中位数。但是这个算法的复杂度是O(m+n)。
leetcode给这道题打上的标签是divide and conquer以及binary search。其实我觉得这个标签有点误导性。百思不得其解之后看了关于这道题的discussion,看完之后不禁感叹:我数学没学好。
这道题的解题思路是这样子的:
该方法的核心是将原问题转变成一个寻找第k小数的问题(假设两个原序列升序排列),这样中位数实际上是第(m+n)/2小的数。所以只要解决了第k小数的问题,原问题也得以解决。
首先假设数组A和B的元素个数都大于k/2,我们比较A[k/2-1]和B[k/2-1]两个元素,这两个元素分别表示A的第k/2小的元素和B的第k/2小的元素。这两个元素比较共有三种情况:>、<和=。如果A[k/2-1]<B[k/2-1],这表示A[0]到A[k/2-1]的元素都在A和B合并之后的前k小的元素中。换句话说,A[k/2-1]不可能大于两数组合并之后的第k小值,所以我们可以将其抛弃。然后在两个数组剩下的元素里寻找第k-(k/2)大的元素(通过递归的方式)。
当A[k/2-1]>B[k/2-1]时存在类似的结论。
当A[k/2-1]=B[k/2-1]时,我们已经找到了第k小的数,也即这个相等的元素,我们将其记为m。由于在A和B中分别有k/2-1个元素小于m,所以m即是第k小的数。(这里可能有人会有疑问,如果k为奇数,则m不是中位数。这里是进行了理想化考虑,在实际代码中略有不同,是先求k/2,然后利用k-k/2获得另一个数。)
如何证明以上的规律是数学书里的内容,这里不做太多介绍。
在实际的编码过程中,我们需要考虑如下边界条件:
1. 如果m或者n为0,直接返回A[k-1]或者B[k-1];
2. 如果k为1,那么返回A[0]或者B[0]当中较小的一个;
3. 如果A[k/2-1] == B[k/2-1],那么直接返回两个数中的任何一个。
代码如下:
class Solution
{
public:
double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2)
{
int m = nums1.size();
int n = nums2.size();
//cout << m << " " << n << endl;
if ((m+n)%2 == 0)
{
//cout << "even" << endl;
return (findKthNum(nums1, nums2, (m + n) / 2) + findKthNum(nums1, nums2, (m + n) / 2 + 1))/2;
}
else
{
//cout << "odd" << endl;
return findKthNum(nums1, nums2, (m + n) / 2 + 1);
}
}
int min(int a, int b)
{
if (a < b)
{
return a;
}
else
{
return b;
}
}
double findKthNum(vector<int>& nums1, vector<int>& nums2, int k)
{
if (nums1.size() > nums2.size())
{
return findKthNum(nums2, nums1, k); //注意nums1的size一定要比nums2小,不然会发生溢出错误
}
//当m或者n为0时
if (nums1.size() == 0)
{
if (k <= nums2.size())
{
return nums2[k - 1];
}
else
{
return nums2[nums2.size() - 1];
}
}
else if (nums2.size() == 0)
{
if (k <= nums1.size())
{
return nums1[k - 1];
}
else
{
return nums1[nums1.size() - 1];
}
}
if (k == 1)
{
return min(nums1[0], nums2[0]);
}
int offset_1 = min(k / 2, nums1.size());
int offset_2 = k - offset_1;
vector<int> new_nums1, new_nums2;
int i = 0;
if (nums1[offset_1-1] == nums2[offset_2-1])
{
return nums1[offset_1-1];
}
else if (nums1[offset_1-1] < nums2[offset_2-1])
{
for (vector<int>::iterator it = nums1.begin(); it != nums1.end(); it++)
{
if (i >= offset_1)
{
new_nums1.push_back(*it);
}
i++;
}
return findKthNum(new_nums1, nums2, k - offset_1);
}
else
{
for (vector<int>::iterator it = nums2.begin(); it != nums2.end(); it++)
{
if (i >= offset_2)
{
new_nums2.push_back(*it);
}
i++;
}
return findKthNum(nums1, new_nums2, k - offset_2);
}
}
};