python-leetcode-寻找两个正序数组的中位数

4. 寻找两个正序数组的中位数 - 力扣（LeetCode）

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        def getKthElement(k):
            """ 主要思路：要找到第 k 小的元素，使用二分查找 """
            index1, index2 = 0, 0
            while True:
                # 特殊情况
                if index1 == len(nums1):
                    return nums2[index2 + k - 1]
                if index2 == len(nums2):
                    return nums1[index1 + k - 1]
                if k == 1:
                    return min(nums1[index1], nums2[index2])

                # 正常情况，比较两个数组中的第 k//2 个元素
                newIndex1 = min(index1 + k // 2 - 1, len(nums1) - 1)
                newIndex2 = min(index2 + k // 2 - 1, len(nums2) - 1)
                pivot1, pivot2 = nums1[newIndex1], nums2[newIndex2]

                if pivot1 <= pivot2:
                    k -= (newIndex1 - index1 + 1)
                    index1 = newIndex1 + 1
                else:
                    k -= (newIndex2 - index2 + 1)
                    index2 = newIndex2 + 1

        m, n = len(nums1), len(nums2)
        total_length = m + n

        if total_length % 2 == 1:
            return getKthElement(total_length // 2 + 1)
        else:
            return (getKthElement(total_length // 2) + getKthElement(total_length // 2 + 1)) / 2