一、原题描述：
二、题目分析：
TopK问题，不管是求前K大/前K小/第K大/第K小等，都有4种不错的方法喔：
1. O(N)：用快排变形最最最高效解决TopK问题
2. O(NlogK)：大根堆（前K小）/小根堆（前K大）
3. O(NlogK)：二叉搜索树
4. O(N): 对于数据范围有限的情况例如本题，可以直接计数排序O(N)高效解决～
下面针对本题，求数组中的前K小的数，用上面的4种方法分别实现一遍。
1. 用快排变形最最最高效解决TopK问题 O(N)
注意找前K大/前K小/第K大/第K小，是不需要对整个数组进行O(NlogN)的排序的！因为可以通过快排切分直接O(N)找到第K大的数（比如求中位数就可以用本方法O(N)找到第mid大的数，如果只会先排序再找的话，那啥...基本上就交代了叭( ͡° ͜ʖ ͡°)...）。
因此本题先通过快排切分排好第K小的数，根据快排切分的性质，它左边的K - 1个数都小于等于它，因此它以及它左边的数就是我们要找的前K小的数。
下面代码给出了详细的注释，没啥好啰嗦的，就是快排模版要记牢哈～
class Solution {
public int[] getLeastNumbers(int[] arr, int k) {
if (k == 0 || arr.length == 0) {
return new int[0];
}
// ⚠️注意最后一个参数传入我们要找的下标（第k小的数下标是k-1）
return quickSearch(arr, 0, arr.length - 1, k - 1);
}
private int[] quickSearch(int[] nums, int lo, int hi, int k) {
// 每快排切分1次，找到排序后下标为j的元素，如果j恰好等于k就返回j以及j左边所有的数；
int j = partition(nums, lo, hi);
if (j == k) {
return Arrays.copyOf(nums, j + 1);
}
// 否则根据下标j与k的大小关系来决定继续切分左段还是右段。
return j > k? quickSearch(nums, lo, j - 1, k): quickSearch(nums, j + 1, hi, k);
}
// 快排切分，返回下标j，使得比nums[j]小的数都在j的左边，比nums[j]大的数都在j的右边。
private int partition(int[] nums, int lo, int hi) {
int v = nums[lo];
int i = lo, j = hi + 1;
while (true) {
while (++i <= hi && nums[i] < v);
while (--j >= lo && nums[j] > v);
if (i >= j) {
break;
}
int t = nums[j];
nums[j] = nums[i];
nums[i] = t;
}
nums[lo] = nums[j];
nums[j] = v;
return j;
}
}快排变形时间复杂度分析：
因为我们是要找下标为k的元素，第一次切分的时候需要遍历整个数组(0 ~ n)找到了下标是j的元素，假如k比j小的话，那么我们下次切分只要遍历数组(0~k-1)的元素就行啦，反之如果k比j大的话，那下次切分只要遍历数组(k+1～n)的元素就行啦，总之平均情况下，可以看作每次调用partition遍历的元素数目都是上一次遍历的1/2，因此时间复杂度是N + N/2 + N/4 + ... + N/N = 2N, 因此时间复杂度是O(N)。
2. 大根堆（前K小）/小根堆（前K大） O(NlogK)
用堆虽然时间复杂度比快排变形慢了点，但是因为Java中提供了现成的PriorityQueue（默认小根堆），所以不需要自己写大段的模版代码，因此实现起来最简单，没几行代码，写起来很快～～面试的时候可以先快点写出这个方案
注意本题是求前K小，因此用一个容量为K的大根堆（每次poll出最大的数，那堆中保留的就是前K小啦）。注意不是小根堆嗷！小根堆的话需要把全部的元素都入堆，那是O(NlogN) ，就不是O(NlogK)啦～～
// 保持堆的大小为K，然后遍历数组中的数字，遍历的时候做如下判断：
// 1. 若目前堆的大小小于K，将当前数字放入堆中。
// 2. 否则判断当前数字与大根堆堆顶元素的大小关系，如果当前数字比大根堆堆顶还大(或等于)，这个数就直接跳过；
// 反之如果当前数字比大根堆堆顶小，先poll掉堆顶，再将该数字放入堆中。
class Solution {
public int[] getLeastNumbers(int[] arr, int k) {
if (k == 0 || arr.length == 0) {
return new int[0];
}
// 默认是小根堆，实现大根堆需要重写一下比较器。
Queue<Integer> pq = new PriorityQueue<>((v1, v2) -> v2 - v1);
for (int num: arr) {
if (pq.size() < k) {
pq.offer(num);
} else if (num < pq.peek()) {
pq.poll();
pq.offer(num);
}
}
// 返回堆中的元素
int[] res = new int[pq.size()];
int idx = 0;
for(int num: pq) {
res[idx++] = num;
}
return res;
}
}3. 二叉搜索树也可以O(NlogK)解决TopK问题嗷～
BST相对没有前两种方法辣么热门，但是也很简单，和大根堆的思路差不多～不得不提的是，与前两种方法相比，BST的优势就是求得的前K个数字保证是有序的。
因为有重复的数字，所以用的是TreeMap而不是TreeSet（有的语言的标准库自带TreeMultiset，也是可以的）。TreeMap的key是数字，value是该数字的个数。我们遍历数组中的数字，维护一个数字总个数为K的TreeMap，每遍历一个元素：
1. 若目前map中数字个数小于K，则将map中当前数字对应的个数+1；
2. 否则，判断当前数字与map中最大的数字的大小关系：若当前数字大于等于map中的最大数字，就直接跳过该数字；若当前数字小于map中的最大数字，则将map中当前数字对应的个数+1，并将map中最大数字对应的个数减1.
class Solution {
public int[] getLeastNumbers(int[] arr, int k) {
if (k == 0 || arr.length == 0) {
return new int[0];
}
// TreeMap的key是数字, value是该数字的个数。
// cnt表示当前map总共存了多少个数字。
TreeMap<Integer, Integer> map = new TreeMap<>();
int cnt = 0;
for (int num: arr) {
// 1. 遍历数组，若当前map中的数字个数小于k，则map中当前数字对应个数+1
if (cnt < k) {
map.put(num, map.getOrDefault(num, 0) + 1);
cnt++;
continue;
}
// 2. 否则，取出map中最大的Key（即最大的数字), 判断当前数字与map中最大数字的大小关系：
// 若当前数字比map中最大的数字还大(或等于)，就直接忽略；
// 若当前数字比map中最大的数字小，则将当前数字加入map中，并将map中的最大数字的个数-1。
Map.Entry<Integer, Integer> entry = map.lastEntry();
if (entry.getKey() > num) {
map.put(num, map.getOrDefault(num, 0) + 1);
if (entry.getValue() == 1) {
map.pollLastEntry();
} else {
map.put(entry.getKey(), entry.getValue() - 1);
}
}
}
// 最后返回map中的元素
int[] res = new int[k];
int idx = 0;
for (Map.Entry<Integer, Integer> entry: map.entrySet()) {
int freq = entry.getValue();
while (freq-- > 0) {
res[idx++] = entry.getKey();
}
}
return res;
}
}说到这里了，就额外再说一下吧，Java的标准库中虽然没有TreeMultiset，但是Guava提供了TreeMultiset，因此也可以通过引入Guava包来实现嗷，不是考察重点，简单贴下大概代码。
// ❌leetcode不支持引入第三方包哦
public int[] getLeastNumbers(int[] arr, int k) {
if (k == 0 || arr.length == 0) {
return new int[0];
}
// TreeMultiset 中允许有重复元素，所以就不需要用TreeMap了。
TreeMultiset<Integer> set = TreeMultiset.create();
for(int num: arr) {
// 1. 遍历每个数字，如果set中的数量小于K，则直接将当前数字加入set中。
if (set.size() < k) {
set.add(num);
continue;
}
// 2. 否则判断当前数字与set中最大数字的大小关系：
// 若当前数字大于等于set中的最大数字，则直接跳过该数字；
// 若当前数字小于set中的最大数字，则将当前数字加入set，并将set中最大数字的个数-1。
Multiset.Entry<Integer> lastEntry = set.lastEntry();
if (num < lastEntry.getElement()) {
set.remove(lastEntry.getElement(), 1);
set.add(num);
}
}
// 返回set中的元素
int[] res = new int[k];
int idx = 0;
for(int num: set) {
res[idx++] = num;
}
return res;
}4. 对于数据范围有限的情况例如本题，可以直接计数排序O(N)高效解决～
class Solution {
public int[] getLeastNumbers(int[] arr, int k) {
if (k == 0 || arr.length == 0) {
return new int[0];
}
// 统计每个数字出现的次数
int[] counter = new int[10001];
for (int num: arr) {
counter[num]++;
}
// 根据counter数组从头找出k个数作为返回结果
int[] res = new int[k];
int idx = 0;
for (int num = 0; num < counter.length; num++) {
while (counter[num]-- > 0 && idx < k) {
res[idx++] = num;
}
if (idx == k) {
break;
}
}
return res;
}
}三、类似题目：
215.数组中的第K个最大元素
扫一扫
774
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
