【LeetCode】4. Median of Two Sorted Arrays

从两个已排序的数组中找到中位数，要求时间复杂度为log(m+n);

苦逼…各种被虐…C/C++、Python，三个版本。

There are two sorted arrays nums1 and nums2 of size m and n respectively.

Find the median of the two sorted arrays. The overall run time complexity should be O(log (m+n)).

Example 1:

nums1 = [1, 3]
nums2 = [2]

The median is 2.0

nums1 = [1, 3]

nums2 = [2]

The median is 2.0

Example 2:

nums1 = [1, 2]
nums2 = [3, 4]

The median is (2 + 3)/2 = 2.5

nums1 = [1, 2]

nums2 = [3, 4]

The median is (2 + 3)/2 = 2.5

题目思路：

因为是log(m+n)，自然想到二分。

假设找到第k小元素(中位数则 k = len/2 —暂不纠结正确好不，反正我知道不对)

假设A的长度小于B的长度 – 如果相反就反过来好了。

pa = min(k/2, lenA)

pb = k – pa

这样分，可以让数组恰好每次去掉一半。

在两个数组中根据pa和pb…将数组分为四部分：(用了python的切片表达)

A[:pa],A[pa:],B[:pb],B[pb:]

然后比较A[pa-1] 和 B[pb-1] 的值，决定去掉其中两块。

为什么这么分，这么删除，自己思考 – 而且一定要思考。

python代码如下：

这个是最先写的，被虐的很爽。

方法不够快，但是很简洁。

我觉得可能是因为切片造成速度太慢。考虑要不要像C++一样传递下标。

嗯，暂时不优化，若觉得有优化需求可以留言。

class Solution(object):
    #这里是找到第k小的元素
    def getKth(self, A, B, k):
        lenA = len(A)
        lenB = len(B)
        
        #其中一方长度为0的情况下，直接返回另一个数组的k-1项即可
        if lenA == 0 and lenB == 0:
            return None
        elif lenA == 0 :
            return B[k-1]
        elif lenB == 0 :
            return A[k-1]
        
        #如果要找到第一小的元素，则比较最小的两个值即可
        if k == 1 :
            return min(A[0], B[0])
        
        #解决长度不均问题
        pa = pb = 0
        if lenA <= lenB :
            pa = min(k/2, lenA)
            pb = k - pa
        else :
            pb = min(k/2, lenB)
            pa = k - pb
        
        #print("pa = %d, pb = %d, k = %d" % (pa, pb, k))
        
        #A[pa] < B[pb]的情况下：
        #A[:pa]内必然没有第k小的元素，因为它们太小了
        #B[pb:]内必然没有第k小的元素，因为它们太大了
        #于是可视为去掉了pa个(下标[0,pa-1])小的元素(大的不考虑)
        #我们就只需要在剩下区间内找到第(pb=)k-pa小的元素
        # A[pa] > B[pb]的情况恰好相反
        #若两者相等，则直接返回它们的值
        
        #注意下标要-1
        if A[pa-1] < B[pb-1] :
            return self.getKth(A[pa:], B[:pb], pb)
        elif A[pa-1] > B[pb-1] :
            return self.getKth(A[:pa], B[pb:], pa)
        else :
            return A[pa-1]
    
    def findMedianSortedArrays(self, nums1, nums2):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :rtype: float
        """
        #获取数组长度
        length1 = len(nums1)
        length2 = len(nums2)
        
        length = (length1 + length2)
        #print("len = %s" % length)
        if length & 1 :
            return self.getKth(nums1, nums2, (length+1)/2)
        else :
            result1 = self.getKth(nums1, nums2, length/2)
            #print("result1 = %d" % (result1))
            
            result2 = self.getKth(nums1, nums2, length/2 + 1)
            #print("result2 = %d" % (result2))
            
            return (result1 + result2) * 0.5

class Solution(object):

#这里是找到第k小的元素

def getKth(self, A, B, k):

lenA = len(A)

lenB = len(B)

#其中一方长度为0的情况下，直接返回另一个数组的k-1项即可

if lenA == 0 and lenB == 0:

return None

elif lenA == 0 :

return B[k-1]

elif lenB == 0 :

return A[k-1]

#如果要找到第一小的元素，则比较最小的两个值即可

if k == 1 :

return min(A[0], B[0])

#解决长度不均问题

pa = pb = 0

if lenA <= lenB :

pa = min(k/2, lenA)

pb = k - pa

else :

pb = min(k/2, lenB)

pa = k - pb

#print("pa = %d, pb = %d, k = %d" % (pa, pb, k))

#A[pa] < B[pb]的情况下：

#A[:pa]内必然没有第k小的元素，因为它们太小了

#B[pb:]内必然没有第k小的元素，因为它们太大了

#于是可视为去掉了pa个(下标[0,pa-1])小的元素(大的不考虑)

#我们就只需要在剩下区间内找到第(pb=)k-pa小的元素

# A[pa] > B[pb]的情况恰好相反

#若两者相等，则直接返回它们的值

#注意下标要-1

if A[pa-1] < B[pb-1] :

return self.getKth(A[pa:], B[:pb], pb)

elif A[pa-1] > B[pb-1] :

return self.getKth(A[:pa], B[pb:], pa)

else :

return A[pa-1]

def findMedianSortedArrays(self, nums1, nums2):

"""

:type nums1: List[int]

:type nums2: List[int]

:rtype: float

"""

#获取数组长度

length1 = len(nums1)

length2 = len(nums2)

length = (length1 + length2)

#print("len = %s" % length)

if length & 1 :

return self.getKth(nums1, nums2, (length+1)/2)

else :

result1 = self.getKth(nums1, nums2, length/2)

#print("result1 = %d" % (result1))

result2 = self.getKth(nums1, nums2, length/2 + 1)

#print("result2 = %d" % (result2))

return (result1 + result2) * 0.5

给出C++代码如下，我不推荐在这里用vector，真XX麻烦！

所以如果是C++读者，我建议用下面的C代码最好不过了。

class Solution {
public:
    double getKth(vector<int>& A, int beginA, int lenA, vector<int>& B, int beginB, int lenB, int k)
    {
        //cout << "lenA = " << lenA << endl;
        //cout << "lenB = " << lenB << endl;
        //cout << "k = " << k << endl;
        
        if(lenA == 0 && lenB == 0)
        {
            return 0.0;
        }
        else if(lenA == 0)
        {
            return B[beginB + k-1];
        }
        else if(lenB == 0)
        {
            return A[beginA + k-1];
        }
        
        if(k == 1)
        {
            return min(A[beginA], B[beginB]);
        }
        
        int pa = 0;
        int pb = 0;
        if(lenA < lenB)
        {
            pa = min(k/2, lenA);
            pb = k - pa;
        }
        else
        {
            pb = min(k/2, lenB);
            pa = k - pb;           
        }
        
        //cout << "pa = " << pa << endl;
        //cout << "pb = " << pb << endl;
        
        if(A[beginA + pa-1] < B[beginB + pb-1])
        {
            //A[pa:], B[:pb]
            return getKth(A, beginA + pa, lenA-pa, B, beginB, pb, pb);
        }
        else if(A[beginA + pa-1] > B[beginB + pb-1])
        {
            //A[:pa], B[pb:]
            return getKth(A, beginA, pa, B, beginB + pb, lenB-pb, pa);
        }
        else
        {
            return A[beginA + pa-1];
        }
    }

    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) 
    {
        int lenA = nums1.size();
        int lenB = nums2.size();
        
        int len = lenA + lenB;
        
        if(len & 1)
        {
            return getKth(nums1,0,lenA,nums2,0,lenB,(len+1)/2);
        }
        else
        {
            return (getKth(nums1,0,lenA,nums2,0,lenB,len/2) + getKth(nums1,0,lenA,nums2,0,lenB,len/2 + 1)) * 0.5 ;
        }
    }
};

class Solution {

public:

double getKth(vector<int>& A, int beginA, int lenA, vector<int>& B, int beginB, int lenB, int k)

{

//cout << "lenA = " << lenA << endl;

//cout << "lenB = " << lenB << endl;

//cout << "k = " << k << endl;

if(lenA == 0 && lenB == 0)

{

return 0.0;

}

else if(lenA == 0)

{

return B[beginB + k-1];

}

else if(lenB == 0)

{

return A[beginA + k-1];

}

if(k == 1)

{

return min(A[beginA], B[beginB]);

}

int pa = 0;

int pb = 0;

if(lenA < lenB)

{

pa = min(k/2, lenA);

pb = k - pa;

}

else

{

pb = min(k/2, lenB);

pa = k - pb;

}

//cout << "pa = " << pa << endl;

//cout << "pb = " << pb << endl;

if(A[beginA + pa-1] < B[beginB + pb-1])

{

//A[pa:], B[:pb]

return getKth(A, beginA + pa, lenA-pa, B, beginB, pb, pb);

}

else if(A[beginA + pa-1] > B[beginB + pb-1])

{

//A[:pa], B[pb:]

return getKth(A, beginA, pa, B, beginB + pb, lenB-pb, pa);

}

else

{

return A[beginA + pa-1];

}

double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2)

{

int lenA = nums1.size();

int lenB = nums2.size();

int len = lenA + lenB;

if(len & 1)

{

return getKth(nums1,0,lenA,nums2,0,lenB,(len+1)/2);

}

else

{

return (getKth(nums1,0,lenA,nums2,0,lenB,len/2) + getKth(nums1,0,lenA,nums2,0,lenB,len/2 + 1)) * 0.5 ;

}

};

给出C语言代码如下：

int min(int a, int b)
{
    return a > b ? b : a;
}

double getKth(int* A, int lenA, int *B, int lenB, int k)
{
    if(lenA == 0 && lenB == 0)
    {
        return 0.0;
    }
    else if(lenA == 0)
    {
        return B[k-1];
    }
    else if(lenB == 0)
    {
        return A[k-1];
    }
    
    if(k == 1)
    {
        return min(A[0], B[0]);
    }
    
    int pa = 0;
    int pb = 0;
    if(lenA < lenB)
    {
        pa = min(k/2, lenA);
        pb = k - pa;
    }
    else
    {
        pb = min(k/2, lenB);
        pa = k - pb;
    }
    
    if(A[pa-1] < B[pb-1])
    {
        //A[pa:], B[:pb]
        return getKth(A+pa, lenA-pa, B, pb, pb);
    }
    else if(A[pa-1] > B[pb-1])
    {
        //A[:pa], B[pb:]
        return getKth(A, pa, B+pb, lenB-pb, pa);
    }
    else
    {
        return A[pa-1];
    }
    //return 0.0;
}

double findMedianSortedArrays(int* nums1, int nums1Size, int* nums2, int nums2Size) 
{
    int len = nums1Size + nums2Size;
    if(len & 1)
    {
        return getKth(nums1, nums1Size, nums2, nums2Size, (len+1)/2);
    }
    else
    {
        return (getKth(nums1, nums1Size, nums2, nums2Size, len/2) 
                + getKth(nums1, nums1Size, nums2, nums2Size, len/2 + 1)) * 0.5;
    }
}

int min(int a, int b)

{

return a > b ? b : a;

}

double getKth(int* A, int lenA, int *B, int lenB, int k)

{

if(lenA == 0 && lenB == 0)

{

return 0.0;

}

else if(lenA == 0)

{

return B[k-1];

}

else if(lenB == 0)

{

return A[k-1];

}

if(k == 1)

{

return min(A[0], B[0]);

}

int pa = 0;

int pb = 0;

if(lenA < lenB)

{

pa = min(k/2, lenA);

pb = k - pa;

}

else

{

pb = min(k/2, lenB);

pa = k - pb;

}

if(A[pa-1] < B[pb-1])

{

//A[pa:], B[:pb]

return getKth(A+pa, lenA-pa, B, pb, pb);

}

else if(A[pa-1] > B[pb-1])

{

//A[:pa], B[pb:]

return getKth(A, pa, B+pb, lenB-pb, pa);

}

else

{

return A[pa-1];

}

//return 0.0;

}

double findMedianSortedArrays(int* nums1, int nums1Size, int* nums2, int nums2Size)

{

int len = nums1Size + nums2Size;

if(len & 1)

{

return getKth(nums1, nums1Size, nums2, nums2Size, (len+1)/2);

}

else

{

return (getKth(nums1, nums1Size, nums2, nums2Size, len/2)

+ getKth(nums1, nums1Size, nums2, nums2Size, len/2 + 1)) * 0.5;

}

不要问我为啥C/C++代码里咋有python注释。

其他疑问可以留言，感谢阅读。

【LeetCode】4. Median of Two Sorted Arrays

Tagged on: Algorithm C++ LeetCode Python