LeetCode 4: Median of Two Sorted Arrays

Median of Two Sorted Arrays

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(n + m)。代码如下:

class Solution {
public:
double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
vector<int> C;
int pa = 0, pb = 0; // index of nums1 and nums2

while (pa < nums1.size() || pb < nums2..size()) {
if (pa == m) {
C.push_back(nums2[pb++]);
continue;
}
if (pb == n) {
C.push_back(nums1[pa++]);
continue;
}
if (nums1[pa] > nums2[pb])
C.push_back(nums2[pb++]);
else
C.push_back(nums1[pa++]);
}
if ((n + m)&1)
return C[(n+m)/2];
else
return (C[(n+m)/2 - 1] + C[(n+m)/2]) / 2.0;
}
};

因为仅需求出中位数，我们可以使用一个计数器m，记录当前已经找到第m大的元素了。同时我们使用两个指针pA和pB，分别指向nums1和nums2数组的第一个元素。使用类似于merge-sort的原理，如果数组nums1当前元素小，那么pA++，同时m++。如果数组nums2当前元素小，那么pB++，同时m++。这样可以得到一个时间复杂度为 O(m+n)，空间复杂度为 O(1)的算法。

思路二：该问题更一般的提法是求有序数组nums1和nums2有序合并之后第k小的数，要求算法的时间复杂度为O(log(m+n))。假设数组nums1和nums2的元素个数都大于k/2，我们比较nums1[k/2-1]和nums2[k/2-1]：若nums1[k/2-1]<nums2[k/2-1]，则nums1[0]到nums2[k/2-1]的元素都在nums1和nums2合并之后的前k小的元素中，可以将其抛弃；若nums1[k/2-1]>nums2[k/2-1]时存在类似的结论；若nums1[k/2-1]=nums2[k/2-1]，我们已经找到了第k小的数。

可以采用递归的方式实现寻找第k小的数，此外我们还需要考虑几个边界条件：

1. 如果nums1或者nums2为空，则直接返回nums2[k-1]或者nums1[k-1]；

2. 如果k为1，我们只需要返回nums1[0]和nums2[0]中的较小值；

3. 如果nums1[k/2-1]=nums2[k/2-1]，返回其中一个。

实现代码如下：

class Solution {
private:
double findKth(vector<int>::iterator iter1, int m, vector<int>::iterator iter2, int n, int k) {
// 总是假设有 m <= n (m 表示从 iter1 往后还有 m 个元素，n 表示从 iter2 往后还有 n 个元素)
if (m > n) return findKth(iter2, n, iter1, m, k);

// 如果A或者B为空，则直接返回B[k-1]或者A[k-1]
if (m == 0) return *(iter2 + k - 1);

// k == 1, 只需比较第一个元素
if (k == 1) return min(*iter1, *iter2);

// 如果k为1，我们只需要返回A[0]和B[0]中的较小值
int pa = min(k / 2, m), pb = k - pa;
if (*(iter1 + pa - 1) < *(iter2 + pb - 1)) {
// A[k/2-1]<B[k/2-1]，这表示A[0]到A[k/2-1]的元素都在A和B合并之后的前k小的元素中，将其抛弃
return findKth(iter1 + pa, m - pa, iter2, n, k - pa);
}
else if (*(iter1 + pa - 1) > *(iter2 + pb - 1)) {
// A[k/2-1]>B[k/2-1]，这表示B[0]到B[k/2-1]的元素都在A和B合并之后的前k小的元素中，将其抛弃
return findKth(iter1, m, iter2 + pb, n - pb, k - pb);
}
else {
// 如果A[k/2-1]=B[k/2-1]，返回其中一个
return *(iter1 + pa - 1);
}
}

public:
double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
int total = nums1.size() + nums2.size();
if (total & 0x1) {
return findKth(nums1.begin(), nums1.size(), nums2.begin(), nums2.size(), total / 2 + 1);
}
else {
int sum = findKth(nums1.begin(), nums1.size(), nums2.begin(), nums2.size(), total / 2);
sum += findKth(nums1.begin(), nums1.size(), nums2.begin(), nums2.size(), total / 2 + 1);
return sum / 2.0;
}
}
};