算法分析与设计——LeetCode:4. 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)).

Example 1:

nums1 = [1, 3]
nums2 = [2]

The median is 2.0

Example 2:

nums1 = [1, 2]
nums2 = [3, 4]

The median is (2 + 3)/2 = 2.5

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
    }
};

思路

我考虑过每次取两个向量的中间的数，比较大小，从而每次去掉m/2或n/2个数，但是这样算法复杂度是  O(log m*n)而不是 O(log (m+n))，而下面的方法通过递归，寻找两个向量的数中第n小的数，每次都可以去掉(m+n)/2个数，从而符合题目的要求。

代码

#include <iostream>
#include <vector>

using namespace std;

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int size1 = nums1.size(), size2 = nums2.size();
        int n = (size1+size2)/2;
        if ((size1+size2)%2 == 0) {
            return (findNum(nums1, size1, nums2, size2, n)+findNum(nums1, size1, nums2, size2, n+1))/2.0;
        } else {
            return findNum(nums1, size1, nums2, size2, n+1);
        }
    }
private:
    int findNum(vector<int> nums1, int size1, vector<int> nums2, int size2, int n) {//找到两个向量所有数中第n小的数
        if (size1 < size2) {
            return findNum(nums2, size2, nums1, size1, n);//使下面运算的size1始终大于等于size2
        }
        if (size2 == 0) {
            return nums1[n-1];
        }
        if (n == 1) {
            if (nums1[0] < nums2[0]) {
                return nums1[0];
            } else {
                return nums2[0];
            }
        }
        int n1 = n/2+n%2;
        int n2 = n-n1;
        if (n2 > size2) {//size2有可能太小，下面的nums[n2-1]会过界
            n1 += n2-size2;
            n2 = size2;
        }
        if (nums1[n1-1] < nums2[n2-1]) {
            nums1.erase(nums1.begin(), nums1.begin()+n1);//截取掉前n1个数
            return findNum(nums1, size1-n1, nums2, size2, n-n1);
        }
        if (nums1[n1-1] > nums2[n2-1]) {
            nums2.erase(nums2.begin(), nums2.begin()+n2);//姐去掉前n2个数
            return findNum(nums1, size1, nums2, size2-n2, n-n2);
        }
        if (nums1[n1-1] == nums2[n2-1]) {
            return nums1[n1-1];
        }
        return 0;
    }
};

int main() {
    vector<int> nums1, nums2;
    nums2 = {1};
    nums1 = {2, 3, 4, 5, 6};
    cout << Solution().findMedianSortedArrays(nums1, nums2) << endl;
    return 0;
}