Python O(log(min(m,n)) solution


  • 25
    Y

    It's guaranteed to be O(log(min(m,n)) because every time the findKth function cuts the shorter array by half of its size.

    class Solution:
        # @return a float
        def findMedianSortedArrays(self, A, B):
            l=len(A)+len(B)
            return self.findKth(A,B,l//2) if l%2==1 else (self.findKth(A,B,l//2-1)+self.findKth(A,B,l//2))/2.0
                
                
        def findKth(self,A,B,k):
            if len(A)>len(B):
                A,B=B,A
            if not A:
                return B[k]
            if k==len(A)+len(B)-1:
                return max(A[-1],B[-1])
            i=len(A)//2
            j=k-i
            if A[i]>B[j]:
                #Here I assume it is O(1) to get A[:i] and B[j:]. In python, it's not but in cpp it is.
                return self.findKth(A[:i],B[j:],i)
            else:
                return self.findKth(A[i:],B[:j],j)

  • 1
    T

    Great job, your solution is so elegant!

    Here is an iteration version of C implementation of your algorithm:

    double findKthSortedArrays(int* nums1, int nums1Size, int*nums2, int nums2Size, int k) {
        int *a1 = nums1;
        int *a2 = nums2;
        int size1 = nums1Size;
        int size2 = nums2Size;
        int *atmp, itmp;
        
        do {
            if (size1 > size2) {
                atmp = a1;
                a1 = a2;
                a2 = atmp;
                itmp = size1;
                size1 = size2;
                size2 = itmp;
            }
            
            if (size1 == 0)
                return a2[k];
            
            if (size1 + size2 == k+1)
                return a1[size1-1] > a2[size2-1] ? a1[size1-1] : a2[size2-1];
            
            int k1 = size1 / 2;
            int k2 = k - k1;
            
            if (a1[k1] > a2[k2]) {
                k = k1;
                a2 += k2;
                size2 -= k2;
                size1 = k1;
            } else {
                k = k2;
                a1 += k1;
                size1 -= k1;
                size2 = k2;
            }
        } while (true);
    }
    
    double findMedianSortedArrays(int* nums1, int nums1Size, int* nums2, int nums2Size) {
        int total = nums1Size + nums2Size;
        if (total & 0x1 == 1)
            return findKthSortedArrays(nums1, nums1Size, nums2, nums2Size, total/2);
        else
            return (findKthSortedArrays(nums1, nums1Size, nums2, nums2Size, total/2 - 1)
                  + findKthSortedArrays(nums1, nums1Size, nums2, nums2Size, total/2)) / 2;
    }
    

  • 0
    S

    Bug found! How can this pass the test?
    "j=k-i"
    j can be out of range of list B.
    Try with
    l1 = [1,3,5]
    l2 = [2,4,6,8]
    k = 5
    You'll see.


  • 0
    S

    Bug found! How can this pass the test?
    "j=k-i"
    j can be out of range of list B.
    Try with
    l1 = [1,3,5]
    l2 = [2,4,6,8] k = 5 You'll see.


  • 1
    Y

    Yes, you can't use the findKth() method to find the kth element in two lists, because j can be out of range. But for this problem, the initialized k is (len(A)+len(B))/2, and if you notice how k is changed during the recursion, you will find that j will be len(B) at most when len(A)==1. When len(A)==1 and k=len(B), len(A)+len(B)-1==k will be matched, the function will return max(A[-1],B[-1]).


  • 1
    L

    It's easy to understand. Learn a lot.
    Perhaps add index as the argument can avoid slicing but will damage the readability.


  • 0

    @yang.zheng.904 you just have to substitute i = min(k // 2, m - 1) and it will work.


  • 1

    @yang-zheng-904 Your solution is brilliant. I modified the kth function a little bit so that now it correctly computes the kth smallest element.

     def kth(self, A, B, k):
            if len(A) > len(B):
                A, B = (B, A)
            if not A:
                return B[k]
            if k == len(A) + len(B) - 1:
                return max(A[-1], B[-1])
        
            i = min(len(A)-1, k/2)
            j = min(len(B)-1, k-i)
        
            if A[i] > B[j]:
                return self.kth(A[:i], B[j:], i)
            else:
                return self.kth(A[i:], B[:j], j)
    

  • 0
    R

    Nice solution!

    I have a question about how to split A and B

    for example

       if A[i]>B[j]:
    
            return self.findKth(A[:i],B[j:],i)
        else:
            return self.findKth(A[i:],B[:j],j)
    

    in the if block the A substation A[:i] will not contains A[i] how you know K won't be A[i], same question for B[:j] in the else block

    thank you very much


  • 0
    W

    @yang.zheng.904 said in Python O(log(min(m,n)) solution:

    return self.findKth(A[:i],B[j:],i)
    Can anyone help me understand why the new k should be i or j?return self.findKth(A[:i],B[j:],i) return self.kth(A[i:], B[:j], j)


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.