面试题30:最小的K个数

题目:输入n个整数,找出其中最小的前K个数。例如输入4,5,1,6,2,7,3,8这8个数字,则最小的4个数字是1,2,3,4.


分析

这是一道很经典的topK问题。前段时间我在参加360实习生面试的时候就考到了这个问题。主要解决方法如下:

  • 排序
  • 基于partition函数
  • 堆(或红黑树)

这三种方法的推荐指数依次递增。

  • 排序

思路很简单,将数组排序后,返回其前K个数即可。由于排序最好的算法时间复杂度为O(nlgn)。所以这种方法的时间复杂度也是O(nlgn)。对于n远大于K的情况,时间效率不是很理想。代码略。

  • partition函数

受益于快速排序的启发,我们可以利用快排中的partition函数来完成。由于一趟partition函数可以确定一个数的最终位置,所以当这个数的下标为K时,返回前面K个数即可。由于返回的前K个数没有排序,所以效率较上一种方法高。因为题目只要求返回前K个数字即可,并没有要求是排序的,所以这种方法是合理的。时间复杂度为O(n)。

代码如下:

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
template<typename T>
void Partition(vector<T>& v,int left,int right,int k)
{
if(left>=right-1)
return;
T tmp=v[left];
int i=left,j=right-1;
while(i<j)
{
while(i<j && v[j]>=tmp)
j--;
swap(v[i],v[j]);
while(i<j && v[i]<=tmp)
i++;
swap(v[i],v[j]);
}
if(i==k)
return;
else if(i<k)
Partition(v,i+1,right,k);
else
Partition(v,left,i,k);
}
template<typename T>
vector<T> GetLeastNumbers(vector<T> v,int k)
{
if(k>v.size())
return v;
Partition(v,0,v.size(),k);
return vector<T>(v.begin(),v.begin()+k);
}

  • 堆或红黑树

这是三种方法中最优的。因为其时间复杂度只有O(nlgK)。在K远小于n时,特别适合于处理海量数据。具体思想如下:因为题目要求找出最小的K个数,我们可以用数组中前K个数建立一个含有K个元素的最大堆,这个堆里面保存的是当前最小的K个数。然后遍历数组中剩下的数字。如果大于堆中的最大值(堆顶元素),说明当前值大于堆中的所有数字,不可能是前K个数字之一。反之如果当钱数字小于堆中最大值,则将最大值pop出去,将当前值插入。由于获取堆顶元素时间复杂度为O(1),插入元素为O(lgK),并且需要遍历一遍长度我n的数组,因此总的时间复杂度为O(nlgK)。

代码如下:

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
template<typename T>
vector<T> GetLeastNumbers(vector<T> v,int k)
{
//堆
if(k>v.size())
return v;
make_heap(v.begin(),v.begin()+k);
int i;
for(i=k;i<v.size();++i)
{
if(v[i]>v[0])
continue;
else
{
swap(v[i],v[0]);
make_heap(v.begin(),v.begin()+k);
}
}
return vector<T>(v.begin(),v.begin()+k);
}
template<typename T>
vector<T> GetLeastNumbers(vector<T> v,int k)
{
//红黑树
if(k>v.size())
return v;
multiset<T> ms;
int i;
for(i=0;i<v.size();++i)
{
if(ms.size()<k)
{
ms.insert(v[i]);
}
else
{
if(v[i]<*ms.begin())
{
ms.erase(ms.begin());
ms.insert(v[i]);
}
}
typename multiset<T>::iterator it=ms.begin();
}
return vector<T>(ms.begin(),ms.end());
}

以上

如果你有任何想法或是可以改进的地方,欢迎和我交流!

完整代码及测试用例在github上:点我前往

<div>来自为知笔记(Wiz)</div>