门| GATE CS 2008 |问题9

题目描述

给定两个有序数组 $a$ 和 $b$，大小分别为 $m$ 和 $n$。找到这两个有序数组中的中位数。要求算法的时间复杂度为 $O(\log(m+n))$。

例如，

$a = [1, 3, 5, 7]$

$b = [2, 4, 6]$

中位数为 $4$。

解题思路

这个问题可以转化为在两个有序数组中找到第 $k$ 小的元素，其中 $k$ 为 $m$ 和 $n$ 的中位数。

可以使用递归二分查找的方法来解决这个问题。比较两个数组的第 $k/2$ 个元素，如果 $a[k/2-1] < b[k/2-1]$，那么 $a[0:k/2]$ 必然不包含中位数。因为在 $a[0:k/2]$ 中最多只能有 $k/2-1$ 个元素，而 $b[0:k/2-1]$ 中已经有 $k/2-1$ 个元素小于 $b[k/2-1]$，所以中位数不可能在 $a[0:k/2]$ 中。因此，我们可以递归地在 $a[k/2: m]$ 和 $b[0:n]$ 中查找第 $k-k/2$ 小的元素。

代码如下：

def findMedianSortedArrays(nums1: List[int], nums2: List[int]) -> float:
    m, n = len(nums1), len(nums2)
    k = (m + n) // 2
    if (m + n) % 2 == 0:
        return (findKth(nums1, nums2, k) + findKth(nums1, nums2, k-1)) / 2.0
    else:
        return findKth(nums1, nums2, k)

def findKth(nums1: List[int], nums2: List[int], k: int) -> int:
    # make sure nums1 is shorter than nums2
    if len(nums1) > len(nums2):
        return findKth(nums2, nums1, k)
    
    m, n = len(nums1), len(nums2)
    if m == 0:
        return nums2[k]
    
    if k == 0:
        return min(nums1[0], nums2[0])
    
    i, j = min(m-1, k//2), min(n-1, k//2)
    if nums1[i] > nums2[j]:
        return findKth(nums1, nums2[j+1:], k-j-1)
    else:
        return findKth(nums1[i+1:], nums2, k-i-1)

时间复杂度为 $O(\log(m+n))$，因为每次递归需要将 $m$ 或 $n$ 减半。返回的是浮点数，需要将中位数转成 float 类型。

测试样例

• 如果 $a = [1, 3, 5, 7]$， $b = [2, 4, 6]$，则中位数为 $4$。
• 如果 $a = [1, 3, 5, 7]$，$ b = [2, 4, 6, 8]$，则中位数为 $4.5$。

总结

这是一道经典的算法题，采用递归二分查找时间复杂度可以做到 $O(\log(m+n))$。在解题过程中注意边界条件的处理。