Median of two Sorted Arrays – Python 3 ( Week 11 – 09)

class Solution:
    def findMedianSortedArrays(self, A: List[int], B: List[int]) -> float:
        l = len(A) + len(B)
        if l % 2 == 0:
            return (self.findKth(A, 0, B, 0, l // 2) + self.findKth(A, 0, B, 0, l // 2 + 1)) / 2
        else:
            return self.findKth(A, 0, B, 0, l // 2 + 1)
            
        
    def findKth(self, A, astart, B, bstart, k):
        if astart >= len(A):
            return B[bstart + k - 1]
        if bstart >= len(B):
            return A[astart + k - 1]
            
        if k == 1:
            return min(A[astart], B[bstart])
            
        a = A[astart + k // 2 - 1] if astart + k // 2 - 1 < len(A) else None
        b = B[bstart + k // 2 - 1] if bstart + k // 2 - 1 < len(B) else None
        
        if b is None or (a is not None and a < b):
            return self.findKth(A, astart + k // 2, B, bstart, k - k // 2 )
        return self.findKth(A, astart, B, bstart + k // 2, k - k // 2 )
        

class Solution:
    """
    @param: A: An integer array
    @param: B: An integer array
    @return: a double whose format is *.5 or *.0
    """
    def findMedianSortedArrays(self, A, B):
        # write your code here
        n = len(A) + len(B)
        if n % 2 == 0:
            return (self.findKth(A, B, n // 2) + self.findKth(A, B, n // 2 + 1) ) / 2.0
        return self.findKth(A, B, n // 2 + 1)
      
    #find Kth of A and B    
    def findKth(self, A, B, k):
        n = len(A)
        m = len(B)
        if n == 0:
            return B[k - 1]
        if m == 0:
            return A[k - 1]
        
        #binary search
        
        start = min(A[0], B[0])
        end = max(A[n - 1], B[m - 1])
        #T = end - start
        
        #log(end - start) times
        #O(logT * log(n+m))
        while start + 1 < end:
            mid = start + (end - start) // 2
            num = self.count_small_or_equal(A, mid) + self.count_small_or_equal(B, mid)
            #Time O(log(n+m))
            if num < k:
                start = mid
            else:
                end = mid
        
        if self.count_small_or_equal(A, start) + self.count_small_or_equal(B, start) == k:
            return start
        return end
        
    def count_small_or_equal(self, arr, target):
            
        start = 0
        end = len(arr) - 1
        while start + 1 < end:
            mid = start + (end - start) // 2
            if arr[mid] <= target:
                start = mid
            else:
                end = mid
        
        if arr[start] > target:
            return start
        elif arr[end] > target:
            return end
        return len(arr)