TopK问题

问题

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

解答

参考：https://zhuanlan.zhihu.com/p/76734219

思路1 排序

将n个数排序之后，取出第k大的数即可。

时间复杂度：$O(n\ast log(n))$

分析：明明只需要TopK，却将全局都排序了，这也是这个方法复杂度非常高的原因。那能不能不全局排序，而只局部排序呢？

思路2 局部排序

不再全局排序，只对最大的k个排序。冒泡是一个很常见的排序方法，每冒一个泡，找出最大值，冒k个泡，就得到TopK。

时间复杂度：$O(n\ast k)$

分析：冒泡，将全局排序优化为了局部排序，非TopK的元素是不需要排序的，节省了计算资源。不少朋友会想到，需求是TopK，是不是这最大的k个元素也不需要排序呢？

* 思路3 堆

思路：只找到TopK，不排序TopK。先用前k个元素生成一个小顶堆，这个小顶堆用于存储，当前最大的k个元素。接着，从第k+1个元素开始扫描，和堆顶（堆中最小的元素）比较，如果被扫描的元素大于堆顶，则替换堆顶的元素，并调整堆，以保证堆内的k个元素，总是当前最大的k个元素。直到，扫描完所有n-k个元素，最终堆中的k个元素，就是TopK。

时间复杂度：$O(n\ast log(k))$

画外音：n个元素扫一遍，假设运气很差，每次都入堆调整，调整时间复杂度为堆的高度，即$lg(k)$，故整体时间复杂度是$n\ast lg(k)$。

分析：堆，将冒泡的TopK排序优化为了TopK不排序，节省了计算资源。堆，是求TopK的经典算法，那还有没有更快的方案呢？

思路4 Partition

除了用来进行快速排序，partition 还可以用 $O(N)$ 的平均时间复杂度解决TopK问题。和快排一样，这里也用到了分而治之的思想。首先用 partition 将数组分为两部分，得到分界点下标 pos，然后分三种情况：

• pos == k-1，则找到第 K 大的值，arr[pos]；
• pos > k-1，则第 K 大的值在左边部分的数组。
• pos < k-1，则第 K 大的值在右边部分的数组。

下面给出基于迭代的实现（这里寻找第 k 小的数字）：

class Solution {
public:
    int partition(vector<int>& nums, int low, int high) {
        int pivot = nums[low];
        while (low < high) {
            while (low < high && nums[high] <= pivot) high--;
            nums[low] = nums[high];
            while (low < high && nums[low] >= pivot) low++;
            nums[high] = nums[low];
        }
        nums[low] = pivot;
        return low;
    } 
    int findKthLargest(vector<int>& nums, int k) {
        int n = nums.size();
        int begin = 0, end = n - 1;
        while (begin < end) {
            int pos = partition(nums, begin, end);
            if (pos == k - 1) {
                return nums[pos];
            } else if (pos < k - 1) {
                begin = pos + 1;
            } else {
                end = pos - 1;
            }
        }
        return nums[begin];//数组长度为1
    }
};

时间复杂度：考虑最坏情况下，每次 partition 将数组分为长度为 $N-1$ 和 $1$ 的两部分，然后在长的一边继续寻找第 K 大，此时时间复杂度为 $O(N^2)$。不过如果在开始之前将数组进行随机打乱，那么可以尽量避免最坏情况的出现（随机选择算在是《算法导论》中一个经典的算法，其时间复杂度为O(n)，是一个线性复杂度的方法。）。而在最好情况下，每次将数组均分为长度相同的两半，运行时间 $T(N)=N+T(N/2)$，时间复杂度是 $O(N)$。

练习

1. leetcode——面试题 17.14. 最小K个数