求第 K 个数的问题

一道经典的题目。给一堆乱序的数，如果它们从小到大排好，求第 k 个是多少。假设排列的下标从 1 开始，而非 0 开始。

这个问题如此之简单而熟悉，可它却可以是很多现实问题的某一个子问题的抽象。它本身相关的问题其实就不少，而且还可以不断演进，成为不同复杂程度的问题。

关于这个问题的分析和演进，我们不妨从一左一右两条分支——堆排序或者快排，来分别进行。在不断演化问题的时候，会这两个分支之间跳来跳去，为了尽量清晰的考虑，我采用一种新方法——使用【分支：堆排序】和【分支：快排】来标注。

Java 中快排用 Arrays.sort 就可以了，如果是堆排序需要用到 PriorityQueue。用 Arrays.sort 写起来最简单（这里的参数校验都省掉了）：

public int getKth(int[] nums, int k) {
    int[] numsCopy = new int[nums.length];
    System.arraycopy(nums, 0, numsCopy, 0, nums.length);

    Arrays.sort(numsCopy);
    return numsCopy[k - 1];
}

【分支：堆排序】我拷贝了一下数组，以免对原数组做修改。当然用 PriorityQueue 写起来不麻烦：

public int getKth(int[] nums, int k) {
    Queue<Integer> heap = new PriorityQueue<>();
    for (int i=0; i<nums.length; i++) {
        heap.add(nums[i]);
    }

    int result = 0;
    for (int i=0; i<k; i++)
        result = heap.poll();

    return result;
}

第一个相关问题来了，Arrays.sort 是怎么实现的，复杂度到底是多少？

【分支：快排】我们可以简单地认为 Arrays.sort 是 n*log(n) 的复杂度。事实上 Java 的实现用的不是普通的快排，而是 DualPivotQuicksort，一个优化了的快速排序算法。一般的快排都是使用一个 pivot，每个周期把比它小的元素扔左边，比它大的扔右边。但是 DualPivotQuicksort 使用了两个 pivot，这样原来这堆数就分要分成三份了。在当今多数的计算机条件下，CPU 计算速度原来越快，而原本被忽略的内存地址访问速度却很难有一样幅度的提高，因而显得越来越举足轻重。因此我们不能只考虑排序过程中单纯的 “大小比较” 次数，还需要考虑实际 “地址访问”（即 num[i]）的开销。因为 CPU 的缓存等原因，在不同情形下，实际对地址访问的次数比算法理论上要少。在有意义的实际应用中，DualPivotQuicksort 因为能够在多数情况下减少地址访问次数而最终比原始的快速排序更快。

第二个引申问题来了，只从算法的角度考虑，是否还有优化余地呢？

如果我只需要找到第 k 个，而不关心从 1 到 k-1 之间元素的顺序，也不关心从 k+1 到最大元素之间的顺序，那能不能通过减少这部分的多余比较，来减少一点运算时间开销呢？其实是可以的。和上面一样，根据堆排序和快排两种思路来梳理优化方法。

【分支：堆排序】先考虑堆排序。我可以修改原来的最小堆实现，由最小堆改为固定堆大小的最大堆。每次放入元素以后检查堆的大小，确保保持在 k。

public int getKth(int[] nums, int k) {
    Queue<Integer> heap = new PriorityQueue<Integer>(k, (o1,o2) -> o2-o1);
    for (int i=0; i<nums.length; i++) {
        heap.add(nums[i]);
        if (i > k - 1)
            heap.poll();
    }

    return heap.poll();
}

注意我初始化的时候初始化了一个大小为 k 的堆，而实际上我维护着的大小是 k-1，这其中有一个留了一个大小为 1 的缓冲，是因为我都是先放进元素，再 poll 来调整堆的大小。因此引入这个缓冲以避免不必要的堆大小 grow。

【分支：快排】再考虑快排优化的思路，每个周期内都把比 pivot 小的往左边扔，比 pivot 大的往右边扔，而这样操作一次以后，我可以知道 pivot 在最后序列中的位置。如果正好是 k，那皆大欢喜；如果比 k 大，说明要找的 k 在这个 pivot 的左边，那就再 k 左边继续进行这样的运算；如果比 k 小，那就再 k 右边继续这样的运算。简单来说就是包含两步：

search：找 pivot 的位置，然后根据和 k 的比较进一步递归 pivot 的左边或者是右边的子数组；
partition：把小的数扔 pivot 左边和把大的数扔 pivot 右边的过程。

细化来说，上述第二步这个和 pivot 比较并且往左或者往右扔数的逻辑是：

先把当前最左边的那个数选举出来作为 pivot（选 pivot 的办法有很多，这只是最简单的一个办法），这里的 pivot 变量实际存储的是它的位置（下标），而其值用变量 x 存储；
然后指针 cur 往右走，如果发现有比 pivot 更小的元素，和 pivot 交换一下，这样操作直到最后；
再把最左边那个数和 pivot 最终应该呆的位置上的数交换一下，就使得 pivot 左边的数都小于 pivot 上的数，pivot 右边的数都大于 pivot 上的数了。

public int getKth(int[] nums, int k) {
    int[] numsCopy = new int[nums.length];
    System.arraycopy(nums, 0, numsCopy, 0, nums.length);

    return search(numsCopy, k - 1, 0, nums.length - 1);
}

private int search(int[] nums, int k, int left, int right) {
    if (left >= right)
        return nums[left];

    int idx = partition(nums, left, right);
    if (idx == k)
        return nums[idx];
    if (idx < k)
        return search(nums, k, idx + 1, right);
    else
        return search(nums, k, left, idx - 1);
}

private int partition(int[] nums, int left, int right) {
    int x = nums[left];
    int pivot = left;
    int cur = left + 1;
    while (cur <= right) {
        if (nums[cur] < x) {
            pivot++;
            swap(nums, pivot, cur);
        }

        cur++;
    }

    swap(nums, left, pivot);

    return pivot;

}

private void swap(int[] nums, int left, int right) {
    if (left == right)
        return;

    nums[left] ^= nums[right];
    nums[right] ^= nums[left];
    nums[left] ^= nums[right];
}

【分支：堆排序】看完快排，下面再回到最开始那个堆排序的分支。如果这堆数很多，但是 k 很小，那使用堆为了取第 k 个数，却需要维护一个巨大的堆，多少显得浪费。于是引出了下面这个问题：

能够改进上面堆排序的做法，仅仅维护一个大小为 k 的堆吗？

最初一想发现不行，为什么？因为堆，或者说优先级队列，只有一个出口。换言之，这个最小堆只能每次去 poll 最小值，如果这个堆的大小已经超过了 k，我要是想从中去掉一个肯定不需要的最大值，是没有办法做到的。

但是什么队列有两个出口呢？Deque。可是一般的 Deque 却又不具备堆的特性，那有没有可能将 PriorityQueue 和 Deque 结合起来呢？这样我的问题就解决了。如果需要我自己实现，那我可以分别创建一个最大堆和一个最小堆，分别保存堆元素的引用。需要 poll 最小值的时候就用最小堆来取，然后拿着取出来的值去最大堆执行删除操作；反之，需要 poll 最大值的时候就用最大堆来取，然后拿着取出来的值去最小堆执行删除操作。代码就不贴了，有上面这些基础，这个很容易实现。

~~不过开源已经有这样的东西了，有不同的实现版本，有的就叫做 PriorityDeque；还有一个版本，是大名鼎鼎的 Guava 实现的，叫做 MinMaxPriorityQueue。~~

Updated：上面这个思路不太对，使用最大堆是可以做到 “仅仅维护一个大小为 k 的堆” 的，具体做法前文已有介绍。

到此先暂停一下，分别看一下改进了的堆排序和改进了的快排的复杂度：

堆排序：时间上看，堆的大小为 k，单一堆操作的时间复杂度是 log(k)，遍历一遍数组是 n，因此最终时间复杂度是 O(n*log(k))；空间复杂度为 O(k)。
快排：空间复杂度为 O(1) 我想这个没有问题，但是时间复杂度呢，如果是单纯的快排，目的是让所有元素有序，复杂度是 O(n*log(n))，但是，找到第 k 个，平均的时间复杂度是在 O(n) 的级别，因为可以这样理解：每次选取一个 pivot 以后，就相当于把原有数组给砍掉一半了，因此最后这个复杂度会是一个等比数列的和——n+n/2+n/4+n/8+……，套用等比数列求和公式，得到 n*(1-(1/2)^n)/(1-(1/2))，在 n 很大的时候，它可以视作 2n，因此复杂度就是 O(n)。

因此二者各有优劣。

如果这堆数不是放在一起，而是在若干个数组里呢？

前面说了，如果这堆数只在一个数组里，有两种办法可以排序，如果是在若干个不同的数组里呢？一样可以从快排和堆排序两个思路去分析。

【分支：快排】如果利用上面改进后的快排，一种方法是合并成一个大数组快排，另一种方法是给每个数组快排之后都各自取最大的 k 个，拿出来放到一起继续快排。

【分支：堆排序】但是倘若 k 相对比较小可以接受而某一个数组太大，那么堆排序就更有优势，因为不需要对这个太大的数组排序，堆的大小只和这个 k 有关，和这些数组本身大小无关。具体做法是，如果数组很大，不需要完全有序，只需要用上面的优化了的方法各自整理出容量为 k 的最小堆，从而使得每个数组都有一个最小堆相对应，能够不断 pull 出当时该数组最小的元素。接着这个问题就变成了，如何从若干个 size 为 k 的最小堆中，找出总体来看排第 k 的元素：

先定义这样一个元素。其中 idx 就存放着第几个堆（下标），num 为实际存放的数值：

class Item {
    int num;
    int idx;

    public Item(int num, int idx) {
        this.num = num;
        this.idx = idx;
    }
}

再在主方法中，对整体维护一个最小堆 heap，每次从这个堆中取出一个元素的时候，要观察这个元素是从哪里来的，如果是从 nums[i] 来的，就再从 nums[i] 取一个当时的最小元素补充到这个 heap 中。

public int getKth(Queue<Integer> nums[], int k) {
    Queue<Item> heap = new PriorityQueue<>(nums.length, (o1, o2) -> o1.num - o2.num);
    for (int i=0; i<k; i++)
        heap.add(new Item(nums[i].poll(), i));

    Item item = null;
    for (int i=0; i<k-1; i++) {
        item = heap.poll();
        heap.add(new Item(nums[item.idx].poll(), item.idx));
    }

    return heap.poll().num;
}

这个方法其实还有一个有趣的推广，就是从一维到二维，甚至更高维。

具体来说，如果拿到若干个数组，从中任意取两个数 x 和 y，要求 x+y 的各种组合里面的第 k 个，或者在全为非负数的情况下引入乘法，比如 x*y+2x 的所有组合里面的第 k 个。这样的问题还是可以基于堆来解决，当然，首先要给每个数组各自排序。思路是类似的。

继续，如果这些数在不同的机器上（文件里）呢？

我想这也是个经典问题，这个问题都问烂了。数据量如果放在一台机器上不合适，那么很多人都会想到，可以 map-reduce 啊，每台机器上进行 map 运算都求出最大的 k 个，然后汇总到一台机器上去 reduce 求出最终的第 k 个（如果机器很多，这个汇总过程可以是多级汇总）。

可是，这个回答无意之中假定了一个条件，让问题变得好处理很多。这个条件就是——k 不大。

假如这堆数很多，因此放在若干台机器上，但是如果这个 k 也非常大呢？即便要想把这 k 个数放到一台机器上去找也不可行。

【分支：快排】这时候问题就有点复杂了，也有不同的处理方法。一种办法是，通过某种排序方法（比如基于不断归并的外排序），给每台机器上的数据都排好序，然后从中找一个（猜一个可能为所求的数）数作为 pivot，并且在每台机器上这些有序数里面都明确这个 pivot 的位置。假设 machine[i] 表示第 i 台机器上，pivot 这个数所处的序列，那么把这些 machine[i] 累加起来，得到的数 sum 去和 k 比较：