Median of Two Sorted Arrays

2020-07-01  本文已影响0人  斯文攸归
class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        lenA = len(nums1)
        lenB = len(nums2)
        if (lenA + lenB) % 2 == 1: 
            #如果A+B的长度是奇数,则找A+B第int((lenA + lenB)/2) + 1大的数
            return self.getKth(nums1, nums2, int((lenA + lenB)/2) + 1)
        else:
            #如果A+B的长度是偶数,则找A+B第int((lenA + lenB)/2)大和int((lenA + lenB)/2)+1大的数,取均值
            return (self.getKth(nums1, nums2, int((lenA + lenB)/2)) + self.getKth(nums1, nums2, int((lenA + lenB)/2) + 1)) * 0.5
    def getKth(self, A, B, k):#寻找A+B中第K大的元素,假设len(A)<len(B)
        lenA = len(A); lenB = len(B)
        if lenA > lenB: 
            return self.getKth(B, A, k)
        if lenA == 0: #直接返回B中第k个元素
            return B[k - 1]
        if k == 1: #直接返回A+B中最小的元素
            return min(A[0], B[0])
        pa = min(int(k/2), lenA)
        #因为len(B)>len(A),所以取k/2和len(A)的较小者,目的是为了比较A和B的第pa个元素
        pb = k - pa
        if A[pa - 1] <= B[pb - 1]:
            #如果A的第pa-1个元素小于等于B的第pb-1个元素,则说明A的前pa个元素可以删掉了,A+B中第k大的元素一定在A[pa:]+B中,但是此时只需要找A[pa:]+B中第pb=k-pa大的元素了
            return self.getKth(A[pa:], B, pb)
        else:
            ##如果A的第pa-1个元素大于等于B的第pb-1个元素,同样说明B的前pb个元素可以删掉了,A+B中第k大的元素一定在A+B[pb:]中,但是此时只需要找A+B[pb:]中第pa=k-pb大的元素了
            return self.getKth(A, B[pb:], pa)
上一篇下一篇

猜你喜欢

热点阅读