面試演算法：lg(k)時間查詢兩個排序陣列合並後第k小的元素

對於一個排好序的陣列A，如果我們要查詢第k小的元素，很簡單，只需要訪問A[k-1]即可，該操作的時間複雜度是O(1).假設給你兩個已經排好序的陣列A和B,他們的長度分別是m和n, 如果把A和B合併成一個排序陣列C, 陣列C含有m+n個元素，要求設計一個演算法，在lg(k)的時間內，找出陣列C中第k小的元素。

例如給定陣列：
A = {1, 3, 5, 7, 9}, B={2, 4, 6, 8, 10} , 合併後的陣列 C = {1,2,3,4,5,6,7,8,9,10}
如果k = 7, 那麼返回的元素是7，也就是A[3]。

一般的處理方法是，先把兩個陣列A和B合併成排好序的C,但是這個過程的時間複雜度是O(m+n), 當然我們可以優化一下，當合並時，只要合併的總元素達到k個就可以，然而這個時間複雜度是O(k)，題目要求的時間複雜度是lg(k)，是個比線性時間還要高的對數級複雜度。

這道題目是難度較大的演算法題，如果能在一個小時內給出演算法並寫出毛病不多的程式碼的話，那麼你的水平已經達到了技術經理的水準。

在演算法設計中，一種思維模式叫逆推法，要想獲得一個結果時，我們可以假設，一旦我們獲得了這個結果後，通過這個結果，我們可以推匯出這個結果會附帶著什麼樣的特殊性質，通過該性質，我們可以得到重大線索，進而推匯出獲取該結果的方法。

根據題目，我們要獲得合併後陣列第k小的元素，這意味著我們從合併陣列的前k個最小元素中，找到最大的那個元素，我們就得到了想要的答案。這前k個元素，要不全部來自陣列A, 要不全部來自陣列B, 要不一部分來自陣列A,一部分來自陣列B,如果k的值比某個陣列的所有元素還要大時，那麼前k個最小元素肯定包含陣列A的全部元素，於是要找到C[k-1], 我們只要找到max(A[m-1], B[k - m - 1])就可以了，例如：
A = {1, 2, 3, 4, 5}, B = {6, 7, 8, 9 ,10}, k = 7,
因此C[6] = max(A[4], B[2]) = B[2] = 7;

因此問題的難度在於第三種情況，也就是前k個最小元素一部分來自陣列A, 一部分來自陣列B。我們用逆推思維看看如何處理這種情況。假設前k個元素中，有l個來自陣列A, 有u個來自陣列B, l + u = k.

於是前k個元素的成分有：A[0], A[1]...A[l-1], B[0], B[1]...B[u-1]。從這個情況我們看看能推匯出什麼性質，我們先假設陣列中的元素不重複。

首先我們有 A[l] > B[u-1] , 要不然A[l] < B[u-1], 那麼我們把B[u-1]拿走，用A[l]替換，那麼所得的k個元素仍然滿足條件，這與我們假設B[u-1]屬於前k個元素的集合相矛盾。由於陣列A是排序的，於是有A[x] > B[u-1] 只要x > l - 1。

同時我們有A[l-1] < B[u], 要不然A[l-1] > B[u], 我們可以把A[l-1]從k個元素集合中拿走，用B[u]來替換，最後得到的k個元素集合仍然滿足條件，這與我們假設A[l-1]屬於k個元素的集合相矛盾，由於陣列A是排序的，因此有A[x] < B[u],只要x < l-1.

根據這兩個性質，我們只要通過查詢到 l-1, 那麼我們就可以找到 u - 1, 進而就能找到第k小的元素。我們可以通過在陣列A中，利用上面提到的兩個性質，通過折半查詢來找到 l - 1 的值。

於是演算法的基本步驟如下，如果陣列A的元素個數比k大，那麼我們就在陣列A的前k個元素中做折半查詢，如果陣列A的元素個數比k小，那麼就在整個陣列A中做折半查詢。

先在陣列A中折半，獲取中間元素假設是A[m/2], 如果A[m/2] > B[k - (m/2+1) - 1] (減1是因為陣列下標從0開始, m/2+1 表示A[m/2]前面包括它自己總共有m/2個元素)那麼l肯定落在0和m/2之間, 如果B[k-(m/2+1)-1] > A[m/2+1] , 那麼l肯定落在區間[m/2, m] 之間，確定區間後，在給定區間中繼續使用折半查詢法，一直找到正確的l為止。我們看個具體例項：
A = {1, 3, 5, 7, 9}, B = {2, 4, 6, 8 ,10}, k = 7

首先在A中折半查詢，找到的元素是A[2] = 5, 對應的B[7 - (2+1) - 1] = B[3] = 8, 此時B[3] > A[3], 所以對於的下標l坐落在區間[2, 4], 在區間[2,4]再次折半，於是得到下標3, A[3] = 7, B[7 - (3+1) -1] = B[2] = 6, 因為B[2] < A[4], 而且A[4] < B[3], 因此 l = 3, u = 2, 也就是說合並後前7小的數有4個來自陣列A, 也就是A[0],A[1],A[2],A[3], 有3個來自陣列B, 也就是B[0], B[1], B[2]。第k小的數只要比較A[3]和B[2]，選出最大那個，根據本例，較大的是A[3], 也就是兩陣列合並後，第k小的數是A[3] = 7。

由於演算法只在一個陣列中折半查詢，並且查詢的範圍不超過k,因此整個演算法複雜度是lg(k),下面我們給出演算法的編碼實現：

public class KthElementSearch {
    private int[] sortedArrayA;
    private int[] sortedArrayB;
    
    private int begin = 0;
    private int end = 0;
    private int requestElementCount = 0;
    private int indexA = 0;
    
    public KthElementSearch(int[] sortedA, int[] sortedB, int k) {
        if (k < 0 || sortedA == null || sortedB == null) {
            throw new IllegalArgumentException("illegal argument");
        }
        
        this.sortedArrayA = sortedA;
        this.sortedArrayB = sortedB;
        
        if (sortedA.length > k - 1) {
            end = k - 1;
        } else {
            end = sortedA.length;
        }
        
        this.requestElementCount = k;
        
        findGivenElement();
    }
    
    private void findGivenElement() {       
        int l = 0;
        
        while (begin <= end) {
            l = (begin + end) / 2;
            /*因為陣列下標從0開始，所以計算當下標是m時，表示總共有l+1個，因此a[l]表示有l+1個元素，
             * b[u]表示有u+1個元素，
             * 
             */
            int u = requestElementCount - (l+ 1) - 1;
            
            
            if (u+1 < sortedArrayB.length && sortedArrayA[l] > sortedArrayB[u+1]) {
                end = l - 1;
            }
            else if (l + 1 < sortedArrayA.length && sortedArrayB[u] > sortedArrayA[l+1]) {
                begin = l + 1;
            } else {
                break;
            }
        
        }
        
        indexA = l;
    }
    
    public int getIndexFromFirstArray() {
        return indexA;
    }
    
    public int getIndexFromSecondArray() {
        return requestElementCount - (indexA + 1) - 1;
    }
}

主入口函式的實現如下：

public static void main(String[] args) {
        int[] A = new int[10];
        int[] B = new int[10];
        int[] C = new int[20];
        int max = 20;
        int min = 1;
        Random random = new Random();
        
        for (int i = 0; i < 10; i++) {
            A[i] = random.nextInt(max) % (max - min + 1) + min;
        }
        Arrays.sort(A);
        System.out.print("Array A is: ");
        for(int i = 0; i < A.length; i++) {
            System.out.print(A[i] + " ");
        }
        
        for (int i = 0; i < 10; i++) {
            B[i] = random.nextInt(max) % (max - min + 1) + min;
        }
        Arrays.sort(B);
        System.out.print("\nArray B is: ");
        for(int i = 0; i < B.length; i++) {
            System.out.print(B[i] + " ");
        }
        
        System.arraycopy(A, 0, C, 0, A.length);
        System.arraycopy(B, 0, C, A.length, B.length);
        Arrays.sort(C);
        System.out.print("\nArra C is: ");
        for (int i = 0; i < C.length; i++) {
            System.out.print(C[i] + " ");
        }
        
        int k = 7;
        System.out.println("\nThe " + k + "th element of combined array is:" + C[k-1]);
        
        
        KthElementSearch kElement = new KthElementSearch(A, B, k);
        int indexA = kElement.getIndexFromFirstArray();
        System.out.println("Index of A is " + indexA + " value of element is: " + A[indexA]);
        int indexB = kElement.getIndexFromSecondArray();
        System.out.println("Index of B is " + indexB + " value of element is: " + B[indexB]);
    }

在主函式中先構造兩個排好序的陣列A和B, 兩陣列中的元素值根據隨機數生成，然後把兩陣列合併成陣列C, 並且先輸出第k小的元素。接著構建KthElementSearch的例項，在該類的實現中，函式findGivenElement實現的就是我們前面說的折半查詢法，getIndexFromFirstArray()返回A陣列對應的元素下標，getIndexFromSecondArray()返回B陣列對應的元素下標，上面程式碼執行後，所得結果如下：

Array A is: 3 6 7 9 12 13 14 14 19 20 
Array B is: 1 2 3 13 14 14 14 14 15 18 
Arra C is: 1 2 3 3 6 7 9 12 13 13 14 14 14 14 14 14 15 18 19 20 
The 7th element of combined array is:9
Index of A is 3 value of element is: 9
Index of B is 2 value of element is: 3

程式先建立了兩個排序陣列A,B,並分別列印出他們元素的內容，同時將兩陣列合併成陣列C, 並給出第7小的元素，它的值是9，接著輸出陣列A元素的對應下標是3， 也就是陣列A的前4個元素組成了合併後陣列C前7小元素的一部分，輸出第二個下標3對應的是陣列B, 也就是陣列B的前3個元素對應合併後陣列C前7小元素的一部分，通過資料對比可以發現，我們演算法得到的結論是正確的，合併後前7小的元素是：1 2 3 3 6 7 9，陣列A前4個元素是：3 6 7 9，陣列B前3個元素是：1 2 3。由此第7小的元素是A[3] = 9, 與程式列印的陣列C第7小元素完全吻合。

更詳細的程式碼除錯和講解請參看視訊：
如何進入google,演算法面試技能全面提升指南

更多技術資訊，包括作業系統，編譯器，面試演算法，機器學習，人工智慧，請關照我的公眾號：