LeetCode 4 [Median of two Sorted Arrays]

原题

两个排序的数组A和B分别含有m和n个数，找到两个排序数组的中位数，要求时间复杂度应为O(log (m+n))。

样例
给出数组A =** [1,2,3,4,5,6]** B = [2,3,4,5]，中位数3.5
给出数组A = [1,2,3] B = [4,5]，中位数 3

解题思路

• 根据题目要求的时间复杂度是O(log(m+n))可知，必须要有个类似于binary search一样，每次操作可以扔一半。
• 如果数组A和B一个含有n个元素，那么,
• 如果n是奇数，则中位数是合并数组的第n/2+1个数
• 如果n是偶数，则中位数是合并数组的第n/2个数和n/2+1个数的平均数
• 最后为题转化为如何求两个排序数组的第k大的数，时间复杂度为O(log(m+n))
• 思路：每次找A的k/2的位置和B的k/2的位置的数与相比较
• 如果A[k/2] < B[k/2]，则扔掉A的前k/2个数，因为第k大的数一定不在其中。找第k大的数转化为，在剩下的A数组和B数组中找第k/2大的数
• 如果A[k/2] >= B[k/2]，则扔掉B的前k/2个数，因为第k大的数一定不在其中。找第k大的数转化为，在剩下的B数组和A数组中找第k/2大的数
• 注意：如果要求返回float类型，要/ 2.0而不是/ 2

完整代码

class Solution(object):
    def findMedianSortedArrays(self, nums1, nums2):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :rtype: float
        """
        n = len(nums1) + len(nums2)
        if n % 2 == 1:
            return self.findKthNum(nums1, nums2, n / 2 + 1)
        else:
            return (self.findKthNum(nums1, nums2, n / 2) + self.findKthNum(nums1, nums2, n / 2 + 1)) / 2.0
            
    def findKthNum(self, nums1, nums2, k):
        if len(nums1) == 0:
            return nums2[k - 1]
        elif len(nums2) == 0:
            return nums1[k -1]
        elif k == 1:
            return min(nums1[0], nums2[0])
        
        a = nums1[k / 2 - 1] if k / 2 <= len(nums1) else None
        b = nums2[k / 2 - 1] if k / 2 <= len(nums2) else None
        if b is None or (a is not None and a < b):
            return self.findKthNum(nums1[k / 2:], nums2, k - k / 2)
        else:
            return self.findKthNum(nums1, nums2[k / 2:], k - k / 2)