经典 TopK 问题

什么是 Top K 问题?简单来说就是在一堆数据里面找到前 K 大(当然也可以是前 K 小)的数。

这个问题也是十分经典的算法问题,不论是面试中还是实际开发中,都非常典型。而这个问题其实也有很多种做法,你真的都懂了么?

经典的 Top K 问题有:最大(小) K 个数、前 K 个高频元素、第 K 个最大(小)元素,下面举例子说明:

栗子(以下所有方法均使用这个栗子)

从 data[1, 12]={5,3,7,1,8,2,9,4,7,2,6,6} 这n=12个数中,找出最大的k=5个。

一、排序

最简单也是最容易想到的方法就是排序,利用快排将数组排序,然后取前k个。

时间复杂度:O(nlogn)

分析:明明只需要TopK,却将全局都排序了,这也是这个方法复杂度非常高的原因。那能不能不全局排序,而只局部排序呢?这就引出了第二个优化方法。

二、局部排序

不再全局排序,只对最大的k个排序。

冒泡是一个很常见的排序方法,每冒一个泡,找出最大值,冒k个泡,就得到TopK。

简单实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
/*方法一 局部冒泡排序*/
vector<int> local_bubble_sort(vector<int> &data, int k)
{
vector<int> res;
for (int i = 0; i < k; ++i)
{
for (int j = i; j < data.size(); ++j)
{
if (data[i] < data[j])
{
swap(data[i], data[j]);
}
}
}
res = vector<int>(data.begin(), data.begin() + k);
return res;
}

时间复杂度:O(n*k)

分析:冒泡,将全局排序优化为了局部排序,非TopK的元素是不需要排序的,节省了计算资源。不少朋友会想到,需求是TopK,是不是这最大的k个元素也不需要排序呢?这就引出了第三个优化方法。

三、堆

思路:只找到TopK,不排序TopK

先用前k个元素生成一个小顶堆,这个小顶堆用于存储,当前最大的k个元素。

接着,从第k+1个元素开始扫描,和堆顶(堆中最小的元素)比较,如果被扫描的元素大于堆顶,则替换堆顶的元素,并调整堆,以保证堆内的k个元素,总是当前最大的k个元素。

直到,扫描完所有n-k个元素,最终堆中的k个元素,就是所求的TopK。

简单实现

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
/*方法二 堆*/
vector<int> local_heap_sort(vector<int>&data, int k)
{
priority_queue<int, vector<int>, greater<int>> heap;
vector<int> res;
for (int num : data)
{
if (heap.size() == k)
{
if (!heap.empty() && heap.top() < num)
{
heap.pop();
heap.emplace(num);
}
}
else
{
heap.emplace(num);
}
}
while (!heap.empty())
{
res.emplace_back(heap.top());
heap.pop();
}
return res;
}

时间复杂度:O(n*log(k))

解释:n个元素扫一遍,最坏的情况下,每插入一个元素都要进行堆调整,调整时间复杂度为堆的高度,即lg(k),故整体时间复杂度是n*lg(k)。

分析:堆,将冒泡的TopK排序优化为了TopK不排序,节省了计算资源。堆,是求TopK的经典算法,那还有没有更快的方案呢?

四、随机选择

随机选择算在是《算法导论》中一个经典的算法,其时间复杂度为O(n),是一个线性复杂度的方法。

这个方法并不是所有同学都知道,为了将算法讲透,先聊一些前序知识,一个所有程序员都应该烂熟于胸的经典算法:快速排序

其伪代码是

1
2
3
4
5
6
7
void quick_sort(int[]arr, int low, inthigh)
{
if(low== high) return;
int i = partition(arr, low, high);
quick_sort(arr, low, i-1);
quick_sort(arr, i+1, high);
}

其核心算法思想是:分治法

分治法(Divide&Conquer),把一个大的问题,转化为若干个子问题(Divide),每个子问题“都”解决,大的问题便随之解决(Conquer)。这里的关键词是“都”。从伪代码里可以看到,快速排序递归时,先通过partition把数组分隔为两个部分,两个部分“都”要再次递归。

分治法有一个特例,叫减治法

减治法(Reduce&Conquer),把一个大的问题,转化为若干个子问题(Reduce),这些子问题中“只”解决一个,大的问题便随之解决(Conquer)。这里的关键词是“只”。

二分查找binary_search,BS,是一个典型的运用减治法思想的算法。

其伪代码是

1
2
3
4
5
6
7
8
9
10
11
12
int BS(int[]arr, int low, inthigh, int target)
{
if(low> high)
return -1;
mid= (low+high)/2;
if(arr[mid]== target)
return mid;
if(arr[mid]> target)
return BS(arr, low, mid-1, target);
else
return BS(arr, mid+1, high, target);
}

从伪代码可以看到,二分查找,一个大的问题,可以用一个mid元素,分成左半区,右半区两个子问题。而左右两个子问题,只需要解决其中一个,递归一次,就能够解决二分查找全局的问题。

通过分治法与减治法的描述,可以发现,分治法的复杂度一般来说是大于减治法的:

快速排序:O(n*lg(n))

二分查找:O(lg(n))

话题收回来,快速排序的核心是:

1
i = partition(arr, low, high);

这个partition是干嘛的呢?

顾名思义,partition会把整体分为两个部分。

更具体的,会用数组 arr中 的一个元素(默认是第一个元素 t=arr[low] )为划分依据,将数据 arr[low, high] 划分成左右两个子数组:

  • 左半部分,都比 t 大

  • 右半部分,都比 t 小

  • 中间位置 i 是划分元素

以上述TopK的数组为例,先用第一个元素 t=arr[low] 为划分依据,扫描一遍数组,把数组分成了两个半区:

  • 左半区比 t 大
  • 右半区比 t 小
  • 中间是 t

partition返回的是 t 最终的位置 i 。

很容易知道,partition的时间复杂度是O(n)。

解释:把整个数组扫一遍,比 t 大的放左边,比 t 小的放右边,最后 t 放在中间 N[i]。

partition和TopK问题有什么关系呢?

TopK是希望求出 arr[1,n] 中最大的 k 个数,那如果找到了第 k 大的数,做一次partition,不就一次性找到最大的 k 个数了么?

解释:即partition后左半区的 k 个数。

问题变成了arr[1, n]中找到第 k 大的数。

再回过头来看看第一次partition,划分之后:

1
i = partition(arr, 1, n);

如果 i 大于 k,则说明 arr[i] 左边的元素都大于 k,于是只递归 arr[1, i-1] 里第 k 大的元素即可;

如果 i 小于 k,则说明说明第 k 大的元素在 arr[i] 的右边,于是只递归 arr[i+1, n] 里第 k-i 大的元素即可;

这就是随机选择算法randomized_select,RS,

其伪代码如下

1
2
3
4
5
6
7
8
9
10
11
int RS(arr, low, high, k)
{
if(low== high)
return arr[low];
i= partition(arr, low, high);
temp= i-low; //数组前半部分元素个数
if(temp>=k)
return RS(arr, low, i-1, k); //求前半部分第k大
else
return RS(arr, i+1, high, k-i); //求后半部分第k-i大
}

简单实现

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
46
47
48
int partition(vector<int>&data, int low, int high)
{
int i = low, j = high, pivot = data[low];
while (i < j)
{
while (i < j&&data[j] > pivot)
{
j--;
}
if (i < j)
{
swap(data[i], data[j]);
i++;
}
while (i<j&&data[i]<pivot)
{
i++;
}
if (i < j)
{
swap(data[i], data[j]);
j--;
}
}
return i;
}
/*方法三 随机选择*/
int randomized_select(vector<int>&data, int low, int high, int k)
{
if (low == high)
{
return low;
}
int pos = partition(data, low, high);
int tmp = pos - low + 1;
if (tmp == k)
{
return pos;
}
else if (tmp > k)
{
return randomized_select(data, low, pos - 1, k);
}
else
{
return randomized_select(data, pos + 1, high, k - tmp);
}
}

这是一个典型的减治算法,递归内的两个分支,最终只会执行一个,它的时间复杂度是O(n)。

再次强调一下:

分治法,大问题分解为小问题,小问题都要递归各个分支,例如:快速排序

减治法,大问题分解为小问题,小问题只要递归一个分支,例如:二分查找,随机选择

通过随机选择(randomized_select),找到 arr[1, n] 中第 k 大的数,再进行一次 partition,就能得到 TopK 的结果。

五、总结

TopK,不难;其思路优化过程,不简单:

全局排序,O(n*lg(n))

局部排序,只排序TopK个数,O(n*k)

堆,TopK个数也不排序了,O(n*lg(k))

分治法,每个分支“都要”递归,例如:快速排序,O(n*lg(n))

减治法,“只要”递归一个分支,例如:二分查找O(lg(n)),随机选择O(n)

TopK的另一个解法:随机选择+partition

原文链接https://blog.csdn.net/z50L2O08e2u4afToR9A/article/details/82837278

参考文章

前端进阶算法10:topk问题

从快速排序、堆排序到top k问题

Author: wnxy
Link: https://wnxy.github.io/2021/06/25/Classic_TopK_problem/
Copyright Notice: All articles in this blog are licensed under CC BY-NC-SA 4.0 unless stating additionally.