LeetCode题解–4. Median of Two Sorted Arrays

链接

LeetCode题目：https://leetcode.com/problems/median-of-two-sorted-arrays/

难度：Hard

题目

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)).

简单来说，此题的目的是找出两个排序数组的中位数，时间复杂度限制在O(log (m+n))。

分析

这题的本质是一个找两个已排序数组中第k大元素的问题。
一开始想到的解法是归并排序，但是这样最坏的情况会比较m+n次，不符合题目的时间复杂度O(log(m+n))的要求。
为了获得一个排序数组的第k大元素，我们可以删除掉前面k-1个元素，然而现在有两个排序数组，可以考虑每次删掉k/2个元素。
假设nums1和nums2的元素个数都超过k/2且m<=n,我们比较nums1[k / 2 - 1]和nums2[k / 2 - 1]的大小，如果前者较小，很容易证明nums1的前k/2个元素都不是第k大,可以把它们都删掉；同理后者较小时删掉nums2的前k/2个元素；而两者相等时则nums1[k / 2 - 1]是中位数。代码可以用递归的方式编写。
论文需要考虑两种边界情况：第一是m < n且第k大元素在nums2的情况，按上面的算法会把nums1的元素删光，此时nums2[k / 2 - 1]是第k大元素；第二是k=1的情况，此时num1[0]和nums2[0]的第k大元素。

代码

#include <vector>
#include <iostream>

using namespace std;

class Solution {
public:

    double findMedianSortedArrays(vector<int> &nums1, vector<int> &nums2) {
        int m = (int) nums1.size();
        int n = (int) nums2.size();
        int sum = m + n;
        if (sum % 2) {
            double median = (double) findKthNum(nums1.begin(), m,
                                                nums2.begin(), n, sum / 2 + 1);
            return median;
        } else {
            double median1 = (double) findKthNum(nums1.begin(), m,
                                                 nums2.begin(), n, sum / 2);
            double median2 = (double) findKthNum(nums1.begin(), m,
                                                 nums2.begin(), n, sum / 2 + 1);
            return (median1 + median2) / 2;
        }
    }

    static int findKthNum(vector<int>::iterator v1, int m,
                          vector<int>::iterator v2, int n, int k) {
        if (m > n) return findKthNum(v2, n, v1, m, k);
        if (m == 0) return *(v2 + k - 1);
        if (k == 1) return min(*v1, *v2);

        int pos1 = min(k / 2, m);
        int pos2 = k - pos1;
        if (*(v1 + pos1 - 1) < *(v2 + pos2 - 1)) {
            return findKthNum(v1 + pos1, m - pos1, v2, n, k - pos1);
        } else if (*(v1 + pos1 - 1) > *(v2 + pos2 - 1)) {
            return findKthNum(v1, m, v2 + pos2, n - pos2, k - pos2);
        } else {
            return *(v1 + pos1 - 1);
        }
    }
};