4. Median of Two Sorted Arrays

Loading...LeetCode

寻找两个 array 的 median，其实就是在两个 array 中找第 K 的大数 (K=(A.length + B.length) / 2)，对两个 sorted array 找 Kth，可以利用 merge sort 的思想来做binary search ：

由于是找第k大数(从1开始)，使用二分法则需要比较A[k/2 - 1]和B[k/2 - 1]，并思考这两个元素和第k大元素的关系。

1. A[k/2 - 1] <= B[k/2 - 1] => A和B合并后的第k大数中必包含A[0]~A[k/2 -1]

2. 若 k/2 - 1 超出A的长度，则必取B[0]~B[k/2 - 1]

利用这一点来不停缩小 k 以及搜索区间，当k == 1时就可以直接返回区间内首个元素。

其他就是考虑一些corner case。

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        len
        len3 = len(nums1) + len(nums2)
        if len3%2 != 0:
            return self.findKth(nums1, 0, nums2, 0, len3//2+1)
        else:
            return (self.findKth(nums1, 0, nums2, 0, len3//2) + self.findKth(nums1, 0, nums2, 0, len3//2+1)) / 2
        
    def findKth(self, A, startA, B, startB, k):
        lenA, lenB = len(A), len(B)
        if startA >= lenA:
            return B[startB + k - 1]
        if startB >= lenB:
            return A[startA + k - 1]
        
        if k == 1: 
            return min(A[startA], B[startB])
        
        midA, midB = float("inf"), float("inf")
        if startA + k//2 - 1 < lenA: midA = A[startA + k//2 - 1]
        if startB + k//2 - 1 < lenB: midB = B[startB + k//2 - 1]
            
        if midA > midB:
            return self.findKth(A, startA, B, startB + k//2, k-k//2)
        else:
            return self.findKth(A, startA + k//2, B, startB, k-k//2)