【经典专题】经典中的经典——TopK问题

问题引入

请找出一堆数据中，最小/最大的k个数。

题目描述非常简单，你有多少种思路去实现它呢？

 
 

解法1——朴素排序

首先可以想到一种非常朴素的思路：将数据从小到大进行排序，取其前k个数即可。

不多赘述，直接看代码：

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
    	// 排序
        Arrays.sort(arr);
        // 取前k个数
        int[] res = new int[k];
        for (int i = 0; i < k; i++) {
            res[i] = arr[i];
        }
        return res;
    }
}

分析：时间复杂度 O(nlogn) ，即排序的开销。不难想到一个问题，我们只需要得到最小的k个数，却没要求这k个数、以及其他n-k个数内部也要有序啊——这种思路似乎进行了多余的排序操作。

 
 

解法2——小顶堆

接着，我们想到了一个关键的数据结构——堆（heap）。

我们把所有数据放到小顶堆中，然后弹出k个数不就行了吗？

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>();
        // 所有数据入堆
        for (int n : arr) {
            heap.offer(n);
        }
        // 弹出k个数
        int[] res = new int[k];
        int index = 0;
        while (index < k) {
            res[index++] = heap.poll();
        }
        return res;
    }
}

分析：这种思路其实和朴素排序的思路没什么不同，对于每个元素，入堆出堆的开销都是O(logn)，因此时间复杂度依旧是 O(nlogn) 。这种思路反而更糟糕，因为这还多使用了 O(n) 的堆空间。

 
 

解法3——大顶堆

接下来，才是堆的正确使用方式 >_< ！

我们维护一个大小为k的大顶堆，将数据依次入堆，当堆的大小超过k时，便弹出一个多出的元素；这个弹出的元素是当前堆中的最大值，它永远不可能包含在最小的k个元素之中；最终堆中的k个元素即为所有数据中最小的k个元素。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>((o1, o2) -> o2 - o1);
        for (int n : arr) {
            heap.offer(n);
            if (heap.size() > k) {
                heap.poll();
            }
        }
        int[] res = new int[k];
        int index = 0;
        while (!heap.isEmpty()) {
            res[index++] = heap.poll();
        }
        return res;
    }
}

还可以进行进一步的优化：如果当前元素已经大于等于堆顶元素的话，那么就直接跳过，反正入堆出堆的都会是它。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (arr.length == 0 || k == 0) {
            return new int[0];
        }
        PriorityQueue<Integer> heap = new PriorityQueue<>((o1, o2) -> o2 - o1);
        for (int n : arr) {
            if (heap.size() < k) {
                heap.offer(n);
                continue;
            }
            if (n < heap.peek()) {
                heap.poll();
                heap.offer(n);
            }
        }
        int[] res = new int[k];
        int index = 0;
        while (!heap.isEmpty()) {
            res[index++] = heap.poll();
        }
        return res;
    }
}

分析：相比于小顶堆，大顶堆优化的本质是什么呢？堆的大小固定为k，而无需装入所有n个元素，因此入堆出堆的开销降为 O(logk)，总的时间复杂度为 O(nlogk) 。另外，空间复杂度也由 O(n) 降为 O(k) 。

 
 

解法4——快速排序

这种思路非常巧妙，也需要有着扎实的基础知识，下面跟着思路体会一下。

本题的要求是：找到左右数据中最小的k个数。

也就等价为：将数据分为前后两组，前面的一组数值较小，后面的一组数值较大，但在这两组的内部并不要求有序。

快速排序的思想是：将数据分为前后两组，前面的一组全部小于基准，后面的一组全部大于基准，但在这两组的内部并不要求有序。

受此启发，得到以下的算法思路：

1）对数据进行一次快速排序，最终基准值（left == right）落在的下标位置为 mid ；

2）此时基准值的位置，就是整体排序完成后的最终位置（这需要你对快速排序的理解比较深刻）；

3）如果 k == mid ，则说明 arr[k] 即为第 k+1 小的数字，那么前k个数字即为最小的k个数字；

4）如果 k < mid ，则说明第 k+1 小的数字在左侧数组中，接着递归左侧数组

5）如果 k > mid ，则说明第 k+1 小的数字在右侧数组中，接着递归右侧数组

看代码吧：

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (arr.length == k) {
            return arr;
        }
        return quickSort(arr, 0, arr.length - 1, k);
    }

    private int[] quickSort(int[] arr, int L, int R, int k) {
        int left = L;
        int right = R;
        int temp = arr[left];
        while (left < right) {
            while (left < right && arr[right] >= temp) {
                right--;
            }
            arr[left] = arr[right];
            while (left < right && arr[left] <= temp) {
                left++;
            }
            arr[right] = arr[left];
        }
        arr[left] = temp;
        if (k < left) {
            return quickSort(arr, L, left - 1, k);
        }
        if (k > left) {
            return quickSort(arr, left + 1, R, k);
        }
        return Arrays.copyOf(arr, k);
    }
}

分析：每次都会根据基准的下标位置和k进行比较，并以此为依据进行递归，每次需要排序的部分都会减半，一个等比数列求和即可得到时间复杂度 O(n) 。空间复杂度即为递归深度 O(logn) 。

 
 

分析总结

1）快速排序思想的使用场景：将数据按某个特征分为两部分，一部分在前，一部分在后，但在这两部分的内部不考虑顺序。

2）使用堆的思路，时间复杂度 O(nlogk) ，使用快速排序的思路，时间复杂度O(n) 。

3）快速排序的思路优于堆的思路吗？从时间复杂度上来看的确如此，但是快速排序的思路有着空间上的局限性：堆可以处理以流的形式到来的大量数据，而快速排序则要求先存储下来所有的数据；当内存不够用的时候，堆反而是解决TopK问题的最优解。

 
 
 
 
 
 
 
 
 
 
 
 

$END$