【算法】Median of Two Sorted Arrays

1. 概述

LeetCode4: https://leetcode.com/problems/median-of-two-sorted-arrays/

求两个有序数组的中位数，要求时间复杂度为 O(log(m + n))

2. 思路分析

很显然，这题应该采用二分查找的思路才能满足时间复杂度 O(log(m + n)) 的要求。

首先需要明确中位数的定义，如果两个数组的长度之和为奇数，那么中位数就是最中间的那个值；如果是偶数，那么中位数应该是最中间的两个数的平均值。一个小trick，(m + n + 1) / 2 和 (m + n + 2) / 2 的下标就是两个数组的中间的位置，如果数组长度之和为奇数，那么这两个得出的是同一个数，如果为偶数，那么分别指向中间位置的两个数。

那么求中位数这个问题就变成了求两个数组中的第 (m + n + 1) / 2 大的数和 (m + n + 2) / 2 大的数。那么就需要一个求第K大的函数。

采用二分思路每次淘汰 k / 2 的数

找出nums1中第k/2大的数midval1和nums2中第k/2大的数midVal2，如果

1. midVal1 <= midVal2：那么nums1中的前k/2个数就可以被丢弃了，因为第k大的数不可能在nums1的前k/2个数中；
2. midVal1 > midVal2：那么同样，nums2中的前k/2个数就可以被丢弃了。

当然，在比较的过程中需要考虑是否越界的问题：

1. 如果nums1中数的个数根本就不够 k/2 个，那么就应该保存nums1中的数，而去淘汰nums2中的数（因为第k大的数不可能在nums2的前k/2个数中的，因为nums1中不足k/2个数。那么我们就可以淘汰nums2的前k/2个数），那么我们可以将midVal1设置为INT_MAX以使其比midVal2大，这样就不会淘汰nums1中的数了；
2. nums2同理。

那么会不会出现nums1和nums2都没有k/2个数的情况呢？在这里的情况是不可能的，因为是求的中位数，那么nums1和nums2中至少有一个包含的数大于等于k/2。

3. 代码和测试结果

Java

package cn.pku.edu.algorithm.leetcode.day01;

/**
 * @author allen
 * @date 2022/9/18
 */
public class LeetCode4 {
    public static void main(String[] args) {
        Solution solution = new Solution();
        int[] nums1 = {1, 3}, nums2 = {2, 4};
        double res = solution.findMedianSortedArrays(nums1, nums2);
        System.out.println(res);
    }
}

class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int left = (nums1.length + nums2.length + 1) / 2;
        int right = (nums1.length + nums2.length + 2) / 2;

        if (left != right) {
            double median1 = findKth(nums1, 0, nums2, 0, left);
            double median2 = findKth(nums1, 0, nums2, 0, right);
            return (median1 + median2) / 2;
        } else {
            return findKth(nums1, 0, nums2, 0, left);
        }
    }

    /**
     * find the k-th element among two sorted arrays
     */
    private double findKth(int[] nums1, int i, int[] nums2, int j, int k) {
        if (i >= nums1.length) return nums2[j + k - 1];
        if (j >= nums2.length) return nums1[i + k - 1];
        if (k == 1) return Math.min(nums1[i], nums2[j]);
        // 取第k/2大的数，如果长度不够，则取整数的最大值，防止被去掉
        int midVal1 = (i + k / 2 - 1 < nums1.length)? nums1[i + k / 2 - 1] : Integer.MAX_VALUE;
        int midVal2 = (j + k / 2 - 1 < nums2.length)? nums2[j + k / 2 - 1] : Integer.MAX_VALUE;
        if (midVal1 <= midVal2) {
            return findKth(nums1, i + k / 2, nums2, j, k - k / 2);
        } else {
            return findKth(nums1, i, nums2, j + k / 2, k - k / 2);
        }
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

C++

#include 
#include 

using namespace std;

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int left = (nums1.size() + nums2.size() + 1) / 2;
        int right = (nums1.size() + nums2.size() + 2) / 2;
        if (left == right) {
            return findKth(nums1, 0, nums2, 0, left);
        } else {
            double median1 = findKth(nums1, 0, nums2, 0, left);
            double median2 = findKth(nums1, 0, nums2, 0, right);
            return (median1 + median2) / 2;
        }
    }

    double findKth(vector<int>& nums1, int i, vector<int>& nums2, int j, int k) {
        if (i >= nums1.size()) return nums2[j + k - 1];
        if (j >= nums2.size()) return nums1[i + k - 1];
        if (k == 1) return min(nums1[i], nums2[j]);
        int midVal1 = (i + k / 2 - 1 < nums1.size())? nums1[i + k / 2 - 1] : INT_MAX;
        int midVal2 = (j + k / 2 - 1 < nums2.size())? nums2[j + k / 2 - 1] : INT_MAX;
        if (midVal1 <= midVal2) {
            return findKth(nums1, i + k / 2, nums2, j, k - k / 2);
        } else {
            return findKth(nums1, i, nums2, j + k / 2, k - k / 2);
        }
    }
};

int main() {
    vector<int> nums1 = {1, 3}, nums2 = {2, 4};
    double res = Solution().findMedianSortedArrays(nums1, nums2);
    cout << res << endl;
    return 0;
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39