【LeetCode 热题100】 4. 寻找两个正序数组的中位数的算法思路及python代码

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

算法思路分析

该代码通过 二分排除法 在两个有序数组中高效寻找中位数，核心思路是逐步缩小搜索范围，每次排除约一半的不可能元素，确保时间复杂度为 $O(log(m+n))$。具体步骤如下：

算法步骤

1. 确定中位数位置：
   • 若总长度 m+n 为奇数，中位数为第 k = (m+n+1)//2 小的元素。
   • 若为偶数，中位数为第 k = (m+n)//2 和 k+1 小元素的平均值。
2. 递归排除法寻找第k小元素：
   • 初始化指针： index1 和 index2 分别表示两个数组的当前起始位置。
   • 边界处理：
     • 若一个数组已全部排除，直接取另一数组的第 k 个元素。
     • 若 k=1，返回两数组当前起始元素的最小值。
   • 计算比较点： 取两数组的 k//2-1 位置元素（若越界则取末尾）。
   • 排除较小部分： 比较两比较点，排除较小值所在数组的前 k//2 个元素，并更新 k 和对应数组的起始索引。

关键点

• 二分排除策略： 每次排除约 k/2 个不可能元素，缩小问题规模。
• 越界处理： 通过 min 函数确保比较点不越界，避免无效访问。
• 动态调整k值： 每次排除后，k 减去已排除元素的数量，确保后续搜索正确。
• 奇偶处理： 根据总长度奇偶性，灵活调用查找函数。

复杂度分析

• 时间复杂度：$O(log(m+n))$
  每次循环排除约一半元素，最多执行 log(m+n) 次。

• 空间复杂度：$O(1)$
  仅使用固定数量的变量，无额外空间消耗。

算法代码

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        def getKthElement(k):
            index1, index2 = 0, 0  # 初始化两个数组的起始索引
            while True:
                # 处理边界：若一个数组被完全排除，直接取另一数组的第k个元素
                if index1 == m:
                    return nums2[index2 + k - 1]
                if index2 == n:
                    return nums1[index1 + k - 1]
                # 当k=1时，只需比较当前索引处的最小值
                if k == 1:
                    return min(nums1[index1], nums2[index2])
                
                # 计算新的比较点，防止越界
                newIndex1 = min(index1 + k // 2 - 1, m - 1)
                newIndex2 = min(index2 + k // 2 - 1, n - 1)
                pivot1, pivot2 = nums1[newIndex1], nums2[newIndex2]
                
                # 排除较小值所在数组的前半部分
                if pivot1 <= pivot2:
                    k -= newIndex1 - index1 + 1  # 更新剩余k值
                    index1 = newIndex1 + 1       # 移动起始索引
                else:
                    k -= newIndex2 - index2 + 1
                    index2 = newIndex2 + 1
        
        m, n = len(nums1), len(nums2)
        totalLength = m + n
        # 根据总长度奇偶性决定中位数计算方式
        if totalLength % 2 == 1:
            return getKthElement((totalLength + 1) // 2)
        else:
            return (getKthElement(totalLength // 2) + getKthElement(totalLength // 2 + 1)) / 2