[LeetCode] 4. Median of Two Sorted Arrays

Posted 2021-01-05 codingEskimo

思路：

这道题要求的时间复杂度是O(log(m+n)), 所以想到的肯定是二分法。但是这道题因为有两个array没有办法直接二分。首先，先判断是基数还是偶数，来确定需要查找第K个数。在查找第K个数的时候，主要思路是各分配k/2个数给两个数组， 看每个数组k/2之前的那个数，小的那个可以确定第K个数不在这个里面，下次继续查找的时候去掉这个部分，并且查找第k - k/2个数。 在具体的查找中，会出现几种情况，第一种是一个其中一个array已经为空，这种情况最简单，说明需要找到的第K个数肯定在另外一个数列里，return出来即可。第二种情况是需要查找第一个数，这种情况只要比较两个数列第一个数中较小的那一个。第三种就是普遍情况的时候，即比较两个数组的K/2-1个数。需要注意的是要处理其中一个数组的长度已经小于K/2了的情况。

 

class Solution:
    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.findKth(nums1, nums2, n//2 + 1)
        else:
            smaller = self.findKth(nums1, nums2, n//2)
            bigger = self.findKth(nums1, nums2, n//2 + 1)
            return (smaller + bigger) / 2.0
        
    def findKth(self, nums1, nums2, k):
            if len(nums1) == 0:
                return nums2[k - 1]
            if len(nums2) == 0:
                return nums1[k - 1] 
            if k == 1:
                return min(nums1[0], nums2[0]) 
            
            e1 = nums1[k//2 - 1] if len(nums1) >= k//2 else None
            e2 = nums2[k//2 - 1] if len(nums2) >= k//2 else None
            
            if e2 is None or (e1 is not None and e1 < e2):
                return self.findKth(nums1[k//2:], nums2, k - k//2)
            return self.findKth(nums1, nums2[k//2:], k - k//2)