LeetCode 4: Median of Two Sorted Arrays

最新推荐文章于 2025-11-25 12:02:55 发布

转载最新推荐文章于 2025-11-25 12:02:55 发布 · 262 阅读

文章标签：

#leetcode

leetcode 专栏收录该内容

91 篇文章

订阅专栏

本文探讨了如何在两个已排序的数组中找到中位数，提出了两种算法思路：一种是将两个数组合并并求中位数，时间复杂度为O(n+m)；另一种是通过递归方式找到第k小的元素，时间复杂度为O(log(m+n))。

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;
        }
    }
};