TopK问题(最大/小的k个数)

TopK问题(最大/小的k个数)

首先简单的陈述一下问题,对于一个整数数组,找出其中最小的k个数,例如输入4、5、1、6、2、7、3、8这8个数字,则最小的4个数字是1、2、3、4。

leetcode链接,这篇博客参考了leetcode的官方解答

求最大或者最小k个数方法都是一样的,下面的解法都是针对最小的k个数

1.排序

最简单的思路就是直接排序呀!排好序后直接选取最小的k个数即可,直接上代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
#include <vector>
#include <iostream>
#include <algorithm>
void getLeastNumbers(vector<int>& arr, int k) {
sort(arr.begin(), arr.end());
}
int main() {
ios::sync_with_stdio(false);
vector<int> arr = {3,2,1,4,6,5};
int k = 3;
getLeastNumbers(arr, k);
for(int i; i<k; i++){
cout<<arr[i]<<" ";
}
return 0;
}

输出结果:

1
1 2 3

算法分析

  • 时间复杂度:O(nlogn),其中 n 是数组 arr 的长度。算法的时间复杂度即排序的时间复杂度。
  • 空间复杂度:O(logn),排序所需额外的空间复杂度为O(logn)

2.堆

我们使用STL容器中的priority_queue实现一个大根堆实时维护数组的前k小值(如果需要前k大值,则需要一个小根堆实时维护数组的 前k大值,priority_queue默认为大根堆,加入参数使用小根堆即可)

直接上代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
void getLeastNumbers(vector<int>& arr, int k) {
priority_queue<int>Q;
for (int i = 0; i < k; ++i) Q.push(arr[i]);
for (int i = k; i < (int)arr.size(); ++i) {
if (arr[i] < Q.top()) {
Q.pop();
Q.push(arr[i]);
}
}
for (int i = 0; i < k; ++i) {
cout<<Q.top()<<" ";
Q.pop();
}
}
int main() {
ios::sync_with_stdio(false);
vector<int> arr = {3,2,1,4,6,5};
int k = 3;
getLeastNumbers(arr, k);
return 0;
}

输出结果:

1
3 2 1

实现小根堆,从而维护最大的k个值:

1
priority_queue<int, vector<int>, greater<int>>Q;

算法分析

  • 时间复杂度:O(nlogk),其中 n 是数组 arr 的长度。由于大根堆实时维护前 k 小值,所以插入删除都是 O(logk) 的时间复杂度,最坏情况下数组里 n 个数都会插入,所以一共需要 O(nlogk) 的时间复杂度。

  • 空间复杂度:O(k),因为大根堆里最多 k 个数

3.快排思想

就是借鉴快速排序的思想,快排的每一次划分都会将数组分为两个部分,而我们现在需要找出最小的k个数其实就是将数组分为两部分,其中一部分个数为k且小于分割点

我们定义randomized_selected(arr, l, r, k)来划分数组,[l,r]为划分数组的范围,k为希望小于划分点的个数,我们调用快排的划分函数划分[l,r]部分的数组,假设得到的划分位置的坐标为pos(小于pos位置的值的位于左侧,大于的位于右侧),然后就会有以下情况:

  1. 如果 pos - l + 1 == k,表示 pivot 就是第 k 小的数,直接获取k左侧的值就是最小的k个数。
  2. 如果 pos - l + 1 < k,表示第 k 小的数在 pivot 的右侧,因此递归调用 randomized_selected(arr, pos + 1, r, k - (pos - l + 1))即可
  3. 如果 pos - l + 1 > k,表示第 k 小的数在 pivot 的左侧,递归调用 randomized_selected(arr, l, pos - 1, k)即可。

这样最终就能找到那个分割点。

下面是代码部分:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
#include <vector>
#include <iostream>
#include <algorithm>
#include <time.h>
// 快排划分的过程,守卫放在最右侧
int partition(vector<int>& nums, int l, int r) {
int pivot = nums[r];
int i = l - 1;
for (int j = l; j <= r - 1; ++j) {
if (nums[j] <= pivot) {
i = i + 1;
swap(nums[i], nums[j]);
}
}
swap(nums[i + 1], nums[r]);
return i + 1;
}
// 基于随机的划分
int randomized_partition(vector<int>& nums, int l, int r) {
int i = rand() % (r - l + 1) + l; // 随机选取划分元素
swap(nums[r], nums[i]);
return partition(nums, l, r); // 返回划分的pos
}
void randomized_selected(vector<int>& arr, int l, int r, int k) {
if (l >= r) return;
int pos = randomized_partition(arr, l, r);
int num = pos - l + 1;
if (k == num) return; // 划分位置刚好为k,直接返回
else if (k < num) randomized_selected(arr, l, pos - 1, k); // 否则继续划分
else randomized_selected(arr, pos + 1, r, k - num);
}
void getLeastNumbers(vector<int>& arr, int k) {
srand((unsigned)time(NULL));
randomized_selected(arr, 0, (int)arr.size() - 1, k);
}
int main() {
ios::sync_with_stdio(false);
vector<int> arr = {3,2,1,4,6,5};
int k = 3;
getLeastNumbers(arr, k);
for(int i = 0; i < k; i++){
cout<<arr[i]<<" ";
}
return 0;
}

输出结果:

1
1 2 3

算法分析

  • 时间复杂度:期望为 O(n),最坏情况下的时间复杂度$O(n^2)$ 为情况最差时,每次的划分点都是最大值或最小值,一共需要划分 n - 1 次,而一次划分需要线性的时间复杂度 O(n),所以最坏情况下时间复杂度为 $O(n^2)$
  • 空间复杂度:期望为O(logn),递归调用的期望深度为 O(logn),每层需要的空间为 O(1),只有常数个变量。

4.bfprt算法

就与快排思想上的改进,由于快排思想在最坏的情况下时间复杂度会达到$O(n^2)$,bfprt算法就是基于此做的一些改进,更详细的说就是快排过程中守卫的选择(当守卫为最大/小值时,算法就会退化到$O(n^2)$ ),通过求两次中位数来选取首位,具体细节可以参考这一篇博客,里面也附有实践代码~

5.直接调用库函数

既然你都坚持看到这里,想必前面的方法都学会了吧,哈哈哈,没想到吧,有现成的库函数(当事人表示很淦,之前咋没发现有这函数QAQ。。)

就是在强大的STL库中存在一个神奇的函数nth_element,就是用来找第k小的整数的,十分的方便(但是感觉也没有快的特别夸张,感觉和自己写的快排思想的topK差不多,但是这个方便鸭)

下面的代码简要介绍一下如何使用,可以说是一目了然了^_^

1
2
3
int a[n];
nth_element(a,a+k,a+n); // 将第k小的元素就位
cout<<a[k]<<endl;