题目
输入整数数组 arr ,找出其中最小的 k 个数。例如,输入4、5、1、6、2、7、3、8这8个数字,则最小的4个数字是1、2、3、4。
思路
TopK问题两种经典解法(面试常问)
- 快速选择(快速排序变体)
- 堆(优先队列)
类似LeetCode HOT 100 —— 215.数组中的第K个最大元素
方法一:快速选择
快速排序中有一步很重要的操作是 partition
(划分),从数组中随机选取一个枢纽元素 pivot
,然后原地移动数组中的元素,使得比 pivot
小的元素在 pivot
的左边,比 pivot
大的元素在 pivot
的右边
这个 partition
操作是原地进行的,需要 O(n)
的时间,接下来,快速排序会递归地排序左右两侧的数组。而快速选择(quick select
)算法的不同之处在于,接下来只需要递归地选择一侧的数组。快速选择算法想当于一个“不完全”的快速排序,因为我们只需要知道最小的 k 个数是哪些,并不需要知道它们的顺序。
目的是寻找最小的 k
个数。假设经过一次 partition
操作,枢纽元素位于下标 m
,也就是说,左侧的数组有 m
个元素,是原数组中最小的 m
个数
- 若
k=m
,就找到了最小的k
个数,就是左侧的数组; - 若
k<m
,则最小的k
个数一定都在左侧数组中,只需要对左侧数组递归地parition
即可; - 若
k>m
,则左侧数组中的m
个数都属于最小的k
个数,还需要在右侧数组中寻找最小的k−m
个数,对右侧数组递归地partition
即可
java代码如下:
class Solution {
public int[] getLeastNumbers(int[] arr, int k) {
if (k == 0 || arr.length == 0) {
return new int[0];
}
// 最后一个参数表示要找的是下标为k-1的数
return quickSearch(arr, 0, arr.length - 1, k - 1);
}
private int[] quickSearch(int[] nums, int l, int r, int k) {
// 每快排切分1次,找到排序后下标为j的元素,如果j恰好等于k就返回j以及j左边所有的数;
int j = partition(nums, l, r);
if (j == k) {
return Arrays.copyOf(nums, j + 1);
}
// 否则根据下标j与k的大小关系来决定继续切分左段还是右段。
return j > k ? quickSearch(nums, l, j - 1, k): quickSearch(nums, j + 1, r, k);
}
// 快排切分,返回下标j,使得比nums[j]小的数都在j的左边,比nums[j]大的数都在j的右边。
private int partition(int[] nums, int l, int r) {
int pivot = nums[l];
int i = l, j = r + 1;
while (true) {
while (++i <= r && nums[i] < pivot);
while (--j >= l && nums[j] > pivot);
if (i >= j) {
break;
}
int t = nums[j];
nums[j] = nums[i];
nums[i] = t;
}
nums[l] = nums[j];
nums[j] = pivot;
return j;
}
}
方法二:堆(优先队列)
堆的性质是每次可以找出最大或最小的元素。
可以使用一个大小为 k
的大顶堆(第/前k小用大顶堆),将数组中的元素依次入堆,当堆的大小超过 k
时,便将多出的元素从堆顶弹出。
java代码如下:
// 保持堆的大小为K,然后遍历数组中的数字,遍历的时候做如下判断:
// 1. 若目前堆的大小小于K,将当前数字放入堆中。
// 2. 否则判断当前数字与大根堆堆顶元素的大小关系,如果当前数字比大根堆堆顶还大,这个数就直接跳过;
// 反之如果当前数字比大根堆堆顶小,先poll掉堆顶,再将该数字放入堆中。
class Solution {
public int[] getLeastNumbers(int[] arr, int k) {
if (k == 0 || arr.length == 0) {
return new int[0];
}
// java中优先队列默认是小根堆,实现大根堆需要重写一下比较器。
Queue<Integer> pq = new PriorityQueue<>((v1, v2) -> v2 - v1);//注意,定义大根堆要降序定义(因为大根堆的根节点最大,要保证降序),要区别出 大根堆用于升序排列(因为大根堆每次出的根节点是最大的,所以用于升序排列)
for (int num: arr) {
if (pq.size() < k) {//如果堆的容量小于k
pq.offer(num);//加入堆
} else if (num < pq.peek()) {//如果堆的容量大于等于k,且要加入的元素num比堆顶元素小,则出堆顶元素,加入num
pq.poll();
pq.offer(num);
}
}
// 返回堆中的元素
int[] res = new int[pq.size()];
int idx = 0;
for(int num: pq) {
res[idx++] = num;
}
return res;
}
}
两种方法的优劣(面试常问)
快速选择 相比 堆,快速选择 的时间复杂度和空间复杂度 都优于使用堆的方法
但是使用快速选择有一些局限性:
- 算法需要修改原数组,如果原数组不能修改的话,还需要拷贝一份数组,空间复杂度就上去了
- 算法需要保存所有的数据。如果把数据看成输入流的话,使用堆的方法是来一个处理一个,不需要保存数据,只需要保存 k 个元素的最大堆。而快速选择的方法需要先保存下来所有的数据,再运行算法。当数据量非常大的时候,甚至内存都放不下的时候,就麻烦了。所以当数据量大的时候还是用基于堆的方法比较好