面试算法：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,算法面试技能全面提升指南

更多技术信息，包括操作系统，编译器，面试算法，机器学习，人工智能，请关照我的公众号：