4. 寻找两个正序数组的中位数

发表于 更新于 分类于

1. 简介

链接：https://leetcode-cn.com/problems/median-of-two-sorted-arrays/

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

示例 1：

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

示例 2：

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

示例 3：

1
2
输入：nums1 = [0,0], nums2 = [0,0]
输出：0.00000

示例 4：

1
2
输入：nums1 = [], nums2 = [1]
输出：1.00000

示例 5：

1
2
输入：nums1 = [2], nums2 = []
输出：2.00000

提示：

  • nums1.length == m
  • nums2.length == n
  • 0 <= m <= 1000
  • 0 <= n <= 1000
  • 1 <= m + n <= 2000
  • -106 <= nums1[i], nums2[i] <= 106

进阶：你能设计一个时间复杂度为 O(log (m+n)) 的算法解决此问题吗？

2. 题解

2.1. 合并数组

最简单的做法就是，合并两个数组，然后进行排序，取出中位数即可。

需要注意的是奇数长度和偶数长度的数组，中位数的取法不同。

1
2
3
4
5
6
7
8
9
10
11
class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        nums = nums1 + nums2
        nums.sort()
        if len(nums)%2==1:
            mid = len(nums)//2
            return nums[mid]
        else:
            mid_2 = len(nums)//2
            mid_1 = mid_2 - 1
            return (nums[mid_1] + nums[mid_2])/2

2.2. 利用二分法

题目希望我们以 $O(log(M+N))$ 的时间复杂度解决此题，log 的时间复杂度，加上两个数组已经是正序的，这可以联想到二分法，因此考虑用二分法解决此题。

根据第一种解法，设合并后的数组长度为 L，奇数长度的数组中位数是第 L//2 + 1 个数字，偶数长度的数组中位数是第 L//2 和 L//2 +1，其实相当于寻找第 k 小的数字1，k 是 L//2 或 L//2 + 1。

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
class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m, n = len(nums1), len(nums2)
        k = (m+n) // 2 + 1
        if (m+n)%2==1:
            return findKth(nums1, 0, nums2, 0, k)
        else:
            return (findKth(nums1, 0, nums2, 0, k-1) + findKth(nums1, 0, nums2, 0, k))/2

def findKth(nums1, start1, nums2, start2, k):
    if start1==len(nums1):
        # 排除了 nums1 的所有元素
        return nums2[start2 + k - 1]
    elif start2==len(nums2):
        # 排除了nums2 的所有元素
        return nums1[start1 + k - 1]
    elif k==1:
        return min(nums1[start1], nums2[start2])
    i = start1 + min(len(nums1), k//2) - 1
    j = start2 + min(len(nums2), k//2) - 1
    if nums1[i] < nums2[j]:
        k = k - (i - start1 + 1)
        start1 = i + 1
        return findKth(nums1, start1, nums2, start2, k)
    else:
        k = k - (j - start2 + 1)
        start2 = j + 1
        return findKth(nums1, start1, nums2, start2, k)

参考资料

1. LeetCode 官方题解：寻找两个有序数组的中位数 ↩