基本思路:寻找最大/最小的k个数,用优先队列(堆)做。因为每次堆超出k个数时需要把堆顶,即最大的那个数pop,所以需要维护大顶堆。
class Solution {
public:
vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
priority_queue<int, vector<int>, less<int>> pq;
vector<int> res;
for (auto &i: input) {
pq.push(i);
if (pq.size() > k) pq.pop();
}
while (!pq.empty()) {
res.push_back(pq.top());
pq.pop();
}
return res;
}
};
优化思路:快速排序的思想是选择一个基准元素,每次用O(n)的时间将所有小于基准元素的元素与所有大于基准元素的元素分开, 然后再对两边的元素递归排序。但是我们需要选择的只是k个最小元素,设所有小于基准元素的元素数目为n。
- 如果n=k则正好满足条件
- 如果n>k,只需要在这n个元素中递归选择即可
- 如果n<k,除了这n个元素外还需要再=在大于基准元素的元素中选择k-n个元素,所以只需要在另外一堆元素中递归选择
这样可以做到常数空间复杂度(如果不计算递归消耗的栈),并且由于每个元素至多只会被遍历一次,所以可以做到线性时间复杂度。
class Solution {
public:
void fastSort(vector<int>& input, int k, int l, int r) {
if (l >= r) return;
int base = input[l];
int sp = l + 1;
for (int i = l + 1; i <= r; i++) {
if (input[i] <= base) {
swap(input[sp++], input[i]);
}
}
swap(input[sp - 1], input[l]);
if (sp == k) return;
else if (sp > k) fastSort(input, k, l, sp - 2);
else fastSort(input, k, sp, r);
}
vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
fastSort(input, k, 0, input.size() - 1);
return vector<int>(input.begin(), input.begin() + k);
}
};

京公网安备 11010502036488号