天天看點

找出兩個排好序的數組的中位數

在LeetCode上看到的一道題目:

給定兩個數組大小分别為m和n,排好了序,可能是降序也可能是升序,求兩個數組所有數字的中位數,要求算法複雜度為O(m+n)。這裡的中位數是如下定義的:如果總個數為偶數那麼就取第n/2和n/2+1個數的平均數,例如:

兩個數組分别為:[1,2] 和[1,2]那麼中位數就應該是1,1,2,2的中位數,也就是:1.5

對于這個題目,最簡單的做法自然是将兩個數組在O(m+n)時間内分别整理成升序排列,然後合并兩個數組到一個大的數組C裡面,最後直接求C的中位數即可,這個做法的代碼我就不寫出來了,還是比較容易寫的,但是最大的缺點在于:浪費空間,數組C的使用就導緻記憶體多了m+n的開銷。

我的做法是類似于合并數組A和B的思路,但是空間開銷要少很多。

首先處理m為0,或者n為0的情況。然後用medianIndex記錄比中位數小的數的個數+1。

緊接着處理A和B的降序或者升序問題。

然後依次從最小的數開始計數(使用count記錄),直到找到medianIndex個數為止。

backwardOne和backwardTwo用來記錄目前為止所找到的最大數和倒數第二大的數,這在m+n為偶數時要用到。

double findMedianSortedArrays(int A[], int m, int B[], int n)
{
    if(m==0&&n!=0)
    {
        return n%2 == 1?B[(n+1)/2-1]:(B[n/2]+B[n/2-1])/2.0;
    }
    if(n==0&&m!=0)
    {
        return m%2 == 1?A[(m+1)/2-1]:(A[m/2]+A[m/2-1])/2.0;
    }
    double result = 0.0;
    int medianIndex = 0;
    if((m+n)%2 == 0)
    {
        medianIndex = (m+n)/2 +1;
    }
    else
    {
        medianIndex = (m+n+1)/2;
    }
    int addA = 1,addB=1;
    int startA = 0 ,endA = m;
    int startB = 0 ,endB = n;
    if(m > 1 && n>1)
    {
        if(A[0] > A[1])
        {
            addA = -1;
            startA = m-1;
        endA = 0;
        }
        if(B[0] > B[1])
        {
            addB = -1;
            startB = n-1;
            endB = 0;
        }
    }
    int count=0;
    int i=0,j=0;

    int backwardOne = 0,backwardTwo = 0;
    for(i=startA,j=startB;i!=endA&&j!=endB;)
    {
        if(A[i]<B[j])
        {
            backwardTwo = backwardOne;
            backwardOne = A[i];
            i+=addA;
            count++;
            if(count == medianIndex)
            {

                return (m+n)%2==1?backwardOne:(backwardOne+backwardTwo)/2.0;
            }
        }
        else
        {
            backwardTwo = backwardOne;
            backwardOne = B[j];
            j+=addB;
            count++;
            if(count == medianIndex)
            {
                return (m+n)%2==1?backwardOne:(backwardOne+backwardTwo)/2.0;
            }
        }
    }

    if(count < medianIndex)
    {
        if(i==endA)
        {

            if(addB == 1)
            {
                if(medianIndex - count >= 2)
                {
                    backwardOne = B[medianIndex - m -2];
                }

                return (m+n)%2==1?B[medianIndex - m -1]:(B[medianIndex - m -1]+backwardOne)/2.0;
            }
            else
            {
                if(medianIndex - count >= 2)
                {
                    backwardOne = B[medianIndex - m];
                }
                return (m+n)%2==1?B[n-(medianIndex - m - 1)]:(B[n-(medianIndex - m - 1)]+backwardOne)/2.0;
            }
        }
        else
        {
            if(addA == 1)
            {
                if(medianIndex - count >= 2)
                {
                    backwardOne = A[medianIndex - n -2];
                }
                return (m+n)%2==1?A[medianIndex - n -1]:(A[medianIndex - n -1]+backwardOne)/2.0;
            }
            else
            {
                if(medianIndex - count >= 2)
                {
                    backwardOne = A[medianIndex - m];
                }
                return (m+n)%2==1?A[m-(medianIndex - n - 1)]:(A[m-(medianIndex - n - 1)]+backwardOne)/2.0;
            }
        }
    }
}