My C++ code (binary-partition)


  • 14
    D

    The basic idea is to always compare the median of A and B and drop half of A or B elements based on the comparison results.

     class Solution {
        public:
        // binary-partition search to find the K-the element     
            int findKSortedArrays(int A[], int m, int B[], int n, int K)
            {
                int med_a, med_b;
        
                if(m<=0)
                { // if A is empty, then return the k-th element of B
                    return B[K-1];
                }
                else if(n<=0)
                { // if B is empty, then return the k-th element of A
                    return A[K-1];
                }
                else
                {
                    //get the median element of A and B and do comparison
                    med_a = m/2;
                    med_b = n/2;
                    
                    if(B[med_b] >= A[med_a])
                    { // if the first half of B and the whole A have NO less than K elements, then dropped the second half of B
                        if(med_a + med_b + 1 >= K)
                            return findKSortedArrays(A, m, B, med_b, K);
                        else  // if the first half of B and the whole A have less than K elements, then dropped the first half of A
                            return findKSortedArrays(&A[med_a+1], m-med_a-1, B, n, K-med_a-1);
                        
                    }
                    else
                    {// if the first half of A and the whole B have NO less than K elements, then dropped the second half of A
                        if(med_a + med_b + 1 >= K)
                            return findKSortedArrays(A, med_a, B, n, K);
                        else // if the first half of A and the whole B have less than K elements, then dropped the first half of B
                            return findKSortedArrays(A, m, &B[med_b+1], n-med_b - 1, K-med_b-1);
                    }
                }
            }
        
            double findMedianSortedArrays(int A[], int m, int B[], int n) {
                double result =0;
                int K; // K is the median element index
                
                if(m>0 || n>0)
                {// if at least one of the arraies are non-empty
                    K = (m+n+1)/2;  
                    result = findKSortedArrays(A, m, B, n, K); // find the K-th element by binary-partition search
                    // if m+n is even, then need to find the K+1-th element and the median is the average of the k-th and k+1-th elements 
                    if((m+n+1)%2)
                    {
                        K = (m+n)/2 + 1; 
                        result = (result  + findKSortedArrays(A, m, B, n, K))/2.0;
                    }
                }
                
                return result;
            }
        };
    

    Updated: I added another concise version here, in which a similar idea is used

    class Solution {
    private:
        double findK(vector<int> &nums1, vector<int> &nums2, int s1, int len1, int s2, int len2, int K)
        {
            if(len1>len2) return findK(nums2,  nums1, s2, len2, s1, len1, K);
            if(len1<=0) return nums2[s2 + K - 1];
            if(K==1) return min(nums1[s1], nums2[s2]);
            int mid1, mid2, hL1, hL2;
            hL1 = min(K/2, len1); 
            mid1 = s1 + hL1 -1;
            hL2 = K - hL1; 
            mid2 = s2 + hL2 -1;
            
            if(nums1[mid1] == nums2[mid2]) return nums1[mid1];
            else if(nums1[mid1] < nums2[mid2]) return findK(nums1, nums2, mid1+1, len1-hL1, s2, len2, K - hL1);
            else return findK(nums1, nums2, s1, len1, mid2+1, len2-hL2, K - hL2);
        }
    public:
        double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
           int len1 = nums1.size(), len2=nums2.size();
           if(!len1 && !len2) return 0;
           if( (len1+len2)%2 ) return findK(nums1, nums2, 0, len1, 0, len2, (len1+len2)/2 + 1);
           else return (findK(nums1, nums2, 0, len1, 0, len2, (len1+len2)/2 ) + findK(nums1, nums2, 0, len1, 0, len2, (len1+len2)/2 + 1 )) / 2.0;
        }

  • -1
    W

    I think you must have made a little mistake, when m+n is even, K = (m + n) / 2 + 1 should be K = (m + n + 1) / 2 + 1, am I right?


  • 0
    A

    would you explain if I can just have the starting index and vector as arguments? Delete the length of nums1 and nums2. Is it possible?

    eg:
    double findK(vector<int> &nums1, vector<int> &nums2, int s1, int s2, int K)


Log in to reply
 

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