解法一
最容易想到的方法是先对元素进行排序,然后取出前k个数,总时间复杂度O(n*logN)。你一定注意到了,当k=1时,上面的算法对整个数组都进行了排序,而原题目只要求最小的k个数,并不需要前k个数有序,也不需要后n-k个数有序。
如何避免n-k个数的排序呢?我们需要部分排序算法,选择排序和交换排序都是不错的选择。把n个数中的前k个数排序出来,复杂度是O(n*k)。
哪一个更好呢?O(nlogN)还是O(nk)?这取决于k的大小。
解法二
在快速排序中,我们用到了一个辅助函数partition。如果基于数组的第k个数字来调整,使得比k个数字小的所有数字都位于数组的左边,比第k个数字大的所有数字都位于数组的右边。这样调整之后,位于数组左边的k个数字就是最小的k个数字(这k个数字不一定是排序的)。基于这种思路的参考代码如下:
vector<int> kSmallNums(vector<int> &nums, int k) {
vector<int> kNums;
int start = 0;
int end = nums.size() - 1;
int index = partition(nums, start, end);
while (index != k - 1) {
if (index > k - 1) {
end = index - 1;
} else {
start = index + 1;
}
index = partition(nums, start, end);
}
for (int i = 0; i < k; i++) {
kNums.push_back(nums[i]);
}
return kNums;
}
解法三:
在前面一种算法中,我们运用了二分法的思想。在本算法中,我们同样借助快速排序中的partition函数,但采用的是一种"分而治之"的思想。
假设N个数村粗在S中,我们从数组S中随机找出一个元素X,把数组分为两部分Sa和Sb。Sa中的元素小于等于X,Sb中的元素大于X。这时,有两种可能性:
- Sa中元素的个数小于K,则Sa中所有的数和Sb中最大的k-|Sa|个元素就是数组S中最大的k个数。
- Sa中元素的个数大于或等于k,则需要返回Sa中最大的k个元素。
参考代码如下:
int partition(vector<int> &nums, int start, int end) {
int randomIndex = rand() % nums.size(); //随机选择一个数作为分组标准
swap(nums[0], nums[randomIndex]); //避免特殊数据下的算法退化
int x = nums[0];
while (start < end) {
while (start < end && nums[end] > x) {
end--;
}
if (start < end) {
nums[start] = nums[end];
start++;
}
while (start < end && nums[start] <= x) {
start++;
}
if (start < end) {
nums[end] = nums[start];
end--;
}
}
nums[start] = x;
return start;
}
vector<int> kSmallNums(vector<int> nums, int k) {
if (k <= 0) {
return vector<int>();
}
if (nums.size() <= k) {
return nums;
}
int index = partition(nums, 0, nums.size() - 1);
vector<int> left = kSmallNums(vector<int>(nums.begin(), nums.begin() + index), k);
vector<int> right = kSmallNums(vector<int>(nums.begin() + index, nums.end()), k - left.size());
left.insert(left.end(), right.begin(), right.end());
return left;
}
解法四
我们已经得到了三种解法,不过这三个解法有个共同之处,就是需要对数据访问多次,那么就有一个问题,如果N很大呢,100亿?这个时候数据就不能全部装入内存,所以要求尽可能少地遍历所有数据。不妨设N>K,前K个数中最小K个数是一种退化的情况,所有K个数就是最小的K个数。如果考虑第K+1个数X呢?如果X比最小的K个数中的最大的数Y大,那么最小的K个数还是保持不变。如果X比Y小,那么最小的K个数应该去掉Y,而包含X。
进一步,我们用容量为K的最大堆来存储最小的K个数。最大堆的堆顶元素就是最小K个数中最大的一个。每次考虑一个数X,如果X比堆顶元素Y大,则不需要更新堆。否则,用X替换堆顶元素Y。在X替换堆顶元素Y之后,X可能破坏最大堆的结构,需要更新堆来维持堆的性质。更新过程花费的时间复杂度O(logK)。在STL中set和multiset都是基于红黑树来实现,我们可以直接拿来用。
参考代码:
vector<int> kSmallNums(vector<int> nums, int k) {
multiset<int, greater<int>> s;
for (int i = 0; i < nums.size() && i < k; i++) {
s.insert(nums[i]);
}
for (int i = k; i < nums.size(); i++) {
int num = nums[i];
if (num < *s.begin()) {
s.erase(s.begin());
s.insert(num);
}
}
vector<int> res;
while (!s.empty()) {
res.push_back(*s.begin());
s.erase(s.begin());
}
return res;
}
总结
基于函数Partition的第二种解法平均时间复杂度O(n*logK),第三种解法平均时间复杂度O(n),他们都需要修改输入的数组。最后一种解法没有修改输入数组,我们每次只是从data中读入数字,所有的写操作都是在堆中进行。二是该算法适合海量数据的输入。由于内存的大小是有限的,有可能不能把海量数据一次性全部读入内存。这个时候,可以从辅助内存中每次读入一个数字。由于每种算法各有优缺点,因此因根据不同的场合采用合适的算法。
网友评论