[数据结构与算法]LeetCode(Python实现)—寻找两个有序数组的中位数

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

题目大意

给定两个大小为 m 和 n 的有序数组 nums1 和 nums2。

请你找出这两个有序数组的中位数，并且要求算法的时间复杂度为 O(log(m + n))。

题目假设 nums1 和 nums2 不会同时为空。

解题思路

根据题目要求，时间复杂度为$O(log(m+n))$，因此无法使用归并排序再查找中位数的方法，需要采用二分查找法。又因题目中所给的数组为有序数组，所以可以将问题简化为寻找两个数组中第k小的数。通过比较A[(m + n) / 2 - 1]和 B[(m + n) / 2 - 1]，对数组中的数字进行排除并更新k值。

例如：
m + n = 14 则 k初始为7，比较A数组和B数组中的第三个元素即A[2]和B[2]的大小。
比较之后，值更小的数组去除第三个元素以及前面的元素，并更新k值，此时k = k - 去掉的元素个数，然后重复上述过程。
PS: 遇到相同值时随机选取一边进行删除
最后当k = 1时，中位数即为当前两个数组的最小值min{A[k]，B[k]}

代码

def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        def getKthValue(k):
            # index 为数组的左边界
            index1, index2 = 0, 0
            while True:
                if index1 == m:
                    # 左边界 + K - 1
                    return nums2[index2 + k - 1]
                if index2 == n:
                    return nums1[index1 + 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)
                value1, value2 = nums1[newIndex1], nums2[newIndex2]
                if value1 <= value2:
                    # k 减去'删除掉的元素个数'
                    k -= newIndex1 - index1 + 1
                    # 修改左边界
                    index1 = newIndex1 + 1
                else:
                    k -= newIndex2 - index2 + 1
                    index2 = newIndex2 + 1
        m, n = len(nums1), len(nums2)
        length = m + n
        if length % 2 == 1:
            return getKthValue((length + 1) // 2)
        else:
            return (getKthValue((length // 2)) + getKthValue(length // 2 + 1)) / 2