参考题解:4种解法秒杀 TopK (快排、堆、二叉搜索树 BST、计数排序)
这道题基于排序做十分简单。但是,由于只需要求最小的 k 个数,因此其实不用全部排序,只需要排序最小的 k 个即可。
因此,考虑做法:
1. 快速排序
class Solution {
public int[] getLeastNumbers(int[] arr, int k) {
if (k == 0 || arr.length == 0) {
return new int[0];
}
return quickSort(arr, 0, arr.length - 1, k - 1);
}
public int[] quickSort(int[] nums, int l, int h, int k) {
int j = partition(nums, l, h);
if (j == k) {
return Arrays.copyOf(nums, j + 1);
}
return j > k ? quickSort(nums, l, j - 1, k) : quickSort(nums, j + 1, h, k);
}
public int partition(int[] nums, int l, int h) {
int v = nums[l];
int i = l, j = h + 1;
while (true) {
while (++i <= h && nums[i] < v);
while (--j >= l && nums[j] > v);
if (i >= j) {
break;
}
int t = nums[j];
nums[j] = nums[i];
nums[i] = t;
}
nums[l] = nums[j];
nums[j] = v;
return j;
}
}
2. 堆
由于要求前 k 小,因此选择大根堆,每次把最大的 poll 出,从而让堆中元素最小。
这里是很有趣的逆向思维。一般我们要堆排序求最小,都是用小根堆。但那样的话,复杂度就是 O(nlogn),而不是 O(nlogk) 了。这里的逆向思维很妙。
基于 Java 提供的 PriorityQueue 可以快速实现。
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. 二叉搜索树 BST
二叉搜索树求得的前 k 个值是有序的,因此是适合求 topK 问题的数据结构。
由于存在重复值,因此采用 TreeMap
实现更合适,而不是 TreeSet。
class Solution {
public int[] getLeastNumbers(int[] arr, int k) {
if (k == 0 || arr.length == 0) {
return new int[0];
}
TreeMap<Integer, Integer> map = new TreeMap<>();
int cnt = 0;
for (int num: arr) {
if (cnt < k) {
map.put(num, map.getOrDefault(num, 0) + 1);
cnt++;
continue;
}
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);
}
}
}
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;
}
}
4. 计数排序
当数据量比较小时,可以直接采用计数排序。
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]++;
}
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;
}
}