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

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

• 标签：数组、二分查找、分治算法
• 难度：困难

题目大意

给定两个正序（从小到大排序）数组 nums1、nums2。要求找出并返回这两个正序数组的中位数。

解题思路

单个有序数组的中位数是中间元素位置的元素。如果中间元素位置有两个元素，则为两个元素的平均数。如果是两个有序数组，则可以使用归并排序的方式将两个数组拼接为一个大的有序数组。大的有序数组中间位置的元素，即为中位数。

当然不合并的话，我们只需找到中位数的位置即可。我们用 n1、n2 来表示数组 nums1、nums2 的长度，则合并后的大的有序数组长度为 (n1 + n2)。

同时可以发现：中位数把数组分割成了左右两部分，并且左右两部分元素个数相等。

• 如果 $(n1 + n2)$ 是奇数时，中位数是大的有序数组中第 $\lfloor \frac{(n1 + n2)}{2} \rfloor + 1$ 的元素，单侧元素个数为 $\lfloor \frac{(n1 + n2)}{2} \rfloor + 1$ 个（包含中位数）。
• 如果 $(n1 + n2)$ 是偶数时，中位数是第 $\lfloor \frac{(n1 + n2)}{2} \rfloor$ 的元素和第 $\lfloor \frac{(n1 + n2)}{2} \rfloor + 1$ 的元素的平均值，单侧元素个数为 $\lfloor \frac{(n1 + n2)}{2} \rfloor$ 个。

因为是向下取整，上面两种情况综合可以写为：单侧元素个数为：$\lfloor \frac{(n1 + n2 + 1)}{2} \rfloor$ 个。

我们用 $k$ 来表示 $\lfloor \frac{(n1 + n2 + 1)}{2} \rfloor$ 。现在的问题就变为了：如何在两个有序数组中找到前 k 个元素？

如果我们从 nums1 数组中取出 m1 个元素，那么从 nums2 就需要取出 m2 = k - m1 个元素。

问题就可以进一步转换为：从 nums1 数组中取出 m1 个元素，那么从 nums2 就需要取出 k - m1 个元素，使得 nums1 第 m1 个元素或 nums2 第 k - m1 个元素为中位线位置。

可以通过「二分查找」来找到合适的 m1 位置。

让 left 指向 $nums1$ 的头部位置 0，right 指向 $nums1$ 的尾部位置 n1。

每次取中间位置 $mid$，判断 $nums1$ 第 $mid+1$ 个元素和 $nums2$ 第 $k - (mid+1)$个元素之间的关系，即 $nums1[mid]$ 和 $nums2[k- mid - 1]$ 的关系。

如果 $nums1[mid] < nums2[k- mid - 1]$，则 $nums1$ 中比 $nums1[mid]$ 小的元素有 $mid$ 个，$nums2$ 数组中即便是前 $k- mid - 2$ 个元素都比 $nums1[mid]$ 小，也最多有 $k- mid - 2$ 个元素比 $nums1[mid]$ 小。所以比 $nums1[mid]$ 小的元素最多有 $mid + (k - mid - 2) = k - 2$ 个。所以 $nums1$ 的前 $mid$ 个元素都不可能是第 $k$ 个元素。可以全部排除。

简单可以描述为：

• 如果 $nums1[mid] < nums2[k- mid - 1]$，则 $nums1$ 的前 $mid$ 个元素都不可能是第 $k$ 个元素，可以全部排除。

• 如果 $nums1[mid] > nums2[k- mid - 1]$，则 $nums2$ 的前 $mid$ 个元素都不可能是第 $k$ 个元素，可以全部排除。

• 如果 $nums1[mid] == nums2[k- mid - 1]$，可以按第一种情况处理。

找到 m1 的位置之后，还要根据两个数组长度和 $(n1 + n2)$ 的奇偶性，以及边界条件来计算对应的中位数。

代码

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        n1 = len(nums1)
        n2 = len(nums2)
        if n1 > n2:
            return self.findMedianSortedArrays(nums2, nums1)

        k = (n1 + n2 + 1) // 2
        left = 0
        right = n1
        while left < right:
            mid = left + (right - left) // 2
            if nums1[mid] < nums2[k-1-mid]:
                left = mid + 1
            else:
                right = mid

        m1 = left
        m2 = k - m1

        c1 = max(float('-inf') if m1 <= 0 else nums1[m1-1], float('-inf') if m2 <= 0 else nums2[m2 - 1])
        if (n1 + n2) % 2 == 1:
            return c1

        c2 = min(float('inf') if m1 >= n1 else nums1[m1], float('inf') if m2 >= n2 else nums2[m2])

        return (c1 + c2) / 2