目录

0004:寻找两个正序数组的中位数(★★)

力扣第 4 题

题目

给定两个大小分别为 mn 的正序(从小到大)数组 nums1nums2。请你找出并返回这两个正序数组的 中位数

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

示例 1:

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

示例 2:

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

提示:

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

相似问题:

分析

  • 假设中位数为 x,那么可以用 x 将 nums1 和 nums2 划分为两部分,满足
    • nums1[i-1]<=x<=nums1[i]
    • nums2[j-1]<=x<=nums2[j]
    • i+j=(m+n)//2
  • 反过来,如果找到一组 (i,j),满足
    • nums1[i-1]<=nums2[j]
    • num2[j-1]<=nums1[i]
    • i+j=(m+n)//2 便定位了中位数
  • 注意到 nums1 和 nums2 都递增,那么只需找到第一个 i 满足
    • nums2[(m+n)//2-i-1]<=nums1[i]
  • 即可用二分查找

再考虑边界情况:

  • 为了方便,当 m>n 时,交换 nums1、nums2,不需计算二分查找 i 的边界
  • m+n 为偶数时,要找到两个中位数取平均
  • 注意 i 和 j 在数组最左/右的边界情况

解答

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m, n = len(nums1),len(nums2)
        if m>n:
            return self.findMedianSortedArrays(nums2,nums1)
        k = (m+n)//2
        i = bisect_left(range(m),1,key=lambda i:nums1[i]>=nums2[k-i-1])
        a = max(nums1[i-1] if i else -inf,nums2[k-i-1] if k-i else -inf)
        b = min(nums1[i] if i<m else inf, nums2[k-i] if k-i<n else inf)
        return b if (m+n)%2 else (a+b)/2

时间 $O(log \ min(M,N))$,48 ms