LeetCode | 寻找两个正序数组的中位数

写在前面#

这题应该是 Hot 100 里面最难的一道了。网上有传言戏称：“面试出这题感觉就是不想要你”。

面试官还请手下留情啊 (๑・́ω・̀๑)

题目#

给定两个大小分别为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。请你找出并返回这两个正序数组的中位数。

算法的时间复杂度应该为 O(log (m+n)) 。

示例 1：

1
输入：nums1 = [1,3], nums2 = [2]
2
输出：2.00000
3
解释：合并数组 = [1,2,3] ，中位数 2

示例 2：

1
输入：nums1 = [1,2], nums2 = [3,4]
2
输出：2.50000
3
解释：合并数组 = [1,2,3,4] ，中位数 (2 + 3) / 2 = 2.5

讲解#

这道题最容易想到的方法是把两个数组合并，再找中位数。这样很好理解，但是需要看完所有元素，时间复杂度是 O(m+n)，不满足题目要求。

最优方法是二分查找。不过这里不是二分查找某个具体数字，而是二分查找一个正确的“切分位置”。

先不要看二分，先看中位数本身。

如果合并后的数组是 [1, 2, 3, 4, 5]，可以把它分成左边 [1, 2, 3] 和右边 [4, 5]。总长度是奇数，左边比右边多一个数，中位数就是左边最大的数，也就是 3。

如果合并后的数组是 [1, 2, 3, 4]，可以把它分成左边 [1, 2] 和右边 [3, 4]。总长度是偶数，中位数就是左边最大值和右边最小值的平均数，也就是 (2 + 3) / 2。

所以这道题真正要找的不是合并后的完整数组，而是一个切分方式。这个切分方式要满足两件事：左边的数量是对的，并且左边所有数都小于等于右边所有数。

假设在 nums1 里切 i 个数到左边，在 nums2 里切 j 个数到左边。为了保证左边数量正确，需要让 i + j = (m + n + 1) // 2。所以只要确定了 i，j 就能跟着确定：j = (m + n + 1) // 2 - i。

这就是为什么可以只在一个数组里二分：我们只需要二分 i，然后用总数量关系算出 j。

为了让二分更安全，代码先保证 nums1 是较短的数组。这样 i 的范围更小，j 也更不容易越界。

下面这张图展示的是一次合法切分：

每次切分后，只需要看切口附近四个数：nums1_left、nums1_right、nums2_left、nums2_right。

因为 nums1 和 nums2 本身都是有序的，所以 nums1 左边最大的数就是 nums1_left，右边最小的数就是 nums1_right；nums2 也是同理。

只要 nums1_left <= nums2_right，就说明 nums1 左边不会比 nums2 右边大。只要 nums2_left <= nums1_right，就说明 nums2 左边不会比 nums1 右边大。两个条件都满足，就代表整个左半边都小于等于整个右半边，这个切分就是正确的。

如果 nums1_left > nums2_right，说明 nums1 左边取多了，i 要往左移。否则说明 nums1 左边取少了，i 要往右移。

找到正确切分以后，中位数就只和切口旁边的数有关。如果总长度是奇数，中位数是 max(nums1_left, nums2_left)；如果总长度是偶数，中位数是 max(nums1_left, nums2_left) 和 min(nums1_right, nums2_right) 的平均数。

这里还用了 float("-inf") 和 float("inf") 处理边界。比如 i == 0 时，nums1 左边没有数，就把 nums1_left 看成负无穷；i == m 时，nums1 右边没有数，就把 nums1_right 看成正无穷。这样代码就不用单独为“切在最左边”或“切在最右边”写很多判断。

代码#

1
from typing import List
2

3

4
class Solution:
5
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
6
        if len(nums1) > len(nums2):
7
            return self.findMedianSortedArrays(nums2, nums1)
8

9
        m = len(nums1)
10
        n = len(nums2)
11
        total_left = (m + n + 1) // 2
12

13
        left = 0
14
        right = m
15

16
        while left <= right:
17
            i = (left + right) // 2
18
            j = total_left - i
19

20
            nums1_left = float("-inf") if i == 0 else nums1[i - 1]
21
            nums1_right = float("inf") if i == m else nums1[i]
22
            nums2_left = float("-inf") if j == 0 else nums2[j - 1]
23
            nums2_right = float("inf") if j == n else nums2[j]
24

25
            if nums1_left <= nums2_right and nums2_left <= nums1_right:
26
                if (m + n) % 2 == 1:
27
                    return float(max(nums1_left, nums2_left))
28
                return (max(nums1_left, nums2_left) + min(nums1_right, nums2_right)) / 2
29

30
            if nums1_left > nums2_right:
31
                right = i - 1
32
            else:
33
                left = i + 1
34

35
        return 0.0

1
from src.python._004_Median_of_Two_Sorted_Arrays import Solution
2

3

4
sol = Solution()
5

6
assert sol.findMedianSortedArrays([1, 3], [2]) == 2.0
7
assert sol.findMedianSortedArrays([1, 2], [3, 4]) == 2.5
8
assert sol.findMedianSortedArrays([], [1]) == 1.0
9
assert sol.findMedianSortedArrays([2], []) == 2.0
10
assert sol.findMedianSortedArrays([0, 0], [0, 0]) == 0.0
11

12
print("PASS")

参考资料#

• 力扣原题：https://leetcode.cn/problems/median-of-two-sorted-arrays/
• 力扣题解页：https://leetcode.cn/problems/median-of-two-sorted-arrays/solutions/
• 二分切分讲解参考：https://liweiwei1419.github.io/leetcode-solution-blog/leetcode-problemset/binary-search/0004-median-of-two-sorted-arrays.html