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

题目描述

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

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

思路

最常规的想法是通过归并排序然后直接寻找中位数，但是这样做的时间复杂度是 $O(m+n)$，题目要求是时间复杂度是 $O(\log (m+n))$，因此我们需要用二分的思想去做题目。

那么如何利用二分的思想去寻找两个数组中的中位数呢？

假设这两个数组分别是 A、B，它们的长度分别为 m、n，这道题目本质上是寻找两个数组中的第 k 小的数字，在这里，k=(m+n)/2 或 k=(m+n)/2+1。

很容易想到，如果我们每次比较 A[k/2-1] 和 B[k/2-1]，会出现几种情况：

• 如果 A[k/2-1]<B[k/2-1]，那么比 A[k/2-1] 小的数字最多有 k-2 个，我们可以排除 A[0] 到 A[k/2-1] 之间的所有值；
• 如果 A[k/2-1]<B[k/2-1]，则相应地排除 B[0]-B[k/2-1]；
• 如果 A[k/2-1]=B[k/2-1]，可以按照第一种情况讨论。

在寻找第 k 小值的时候，我们可以用 idx1、idx2 表示当前数组的左起始点。

为了寻找第 k 小的数字，可以寻找新的坐标 newIdx1、newIdx2 为 min(idx1+k//2-1, m-1)、min(idx2+k//2-1, n-1)，这是为了防止数组溢出。

接下来比较对应位置数字的大小，按照上面的想法我们可以更新左边界，并缩小 k 的值。假设 nums1[newIdx1] <= nums2[newIdx2] 则：

• 每次缩小的 k 为 (newIdx1-idx1+1)
• 每次缩小的更新的边界值为 newIdx1+1

依次循环即可。

最终代码：

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        def getKthElem(k):
            idx1, idx2 = 0, 0
            while True:
                # 边界条件，列表空
                if idx1 == m:
                    return nums2[idx2+k-1]
                if idx2 == n:
                    return nums1[idx1+k-1]
                if k == 1:
                    return min(nums1[idx1], nums2[idx2])
                newIdx1 = min(idx1+k//2-1, m-1)
                newIdx2 = min(idx2+k//2-1, n-1)
                if nums1[newIdx1] <= nums2[newIdx2]:
                    k -= (newIdx1-idx1+1)
                    idx1 = newIdx1+1
                else:
                    k -= (newIdx2-idx2+1)
                    idx2 = newIdx2 + 1
        m, n = len(nums1), len(nums2)
        if (m+n) % 2 == 1:
            return getKthElem((m+n+1)//2)
        else:
            return (getKthElem((m+n)//2)+getKthElem((m+n)//2+1))/2

想得太多