0004：寻找两个正序数组的中位数（★★）

力扣第 4 题

题目

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

算法的时间复杂度应该为 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