Problem Source : https://leetcode-cn.com/problems/smallest-k-lcci/
Solution Source : https://github.com/hujingbo98/algorithm/blob/master/source/leetcodeInterviewProblem/17_14_SmallestKNumbers.cpp
面试题 17.14、最小K个数
设计一个算法,找出数组中最小的k个数。以任意顺序返回这k个数均可。
示例:
输入:arr = [1,3,5,7,2,4,6,8], k = 4
输出:[1,2,3,4]
提示:
0 <= len(arr) <= 100000
0 <= k <= min(100000, len(arr))
方法一:最大堆
使用一个最大堆实时维护数组的前 k 小值。首先,先将数组的前 k 个数插入到最大堆中,随后开始从数组的第 k + 1 个数遍历,如果遍历的数比最大堆的堆顶的数小,就将堆顶的数弹出,再将遍历的数插入最大堆。最后将最大堆里的数存入数组返回即可。
时间复杂度是 O((n + k)(log k)),其中 n 是数组的长度。时间复杂度分析:
- 将数组的前 k 个数插入到最大堆的时间是 O(k(log k))
- 从 k + 1,开始遍历,最坏情况是都需要插入到最大堆,此时,时间复杂度是 O((n-k)(log k))
- 将最大堆里的数存入数组的时间复杂度是 O(k(log k))
空间复杂度是 O(k),最大堆里最多存 k 个数。
vector<int> smallestK(vector<int>& arr, int k) {
vector<int> ans;
if (k <= 0) return ans;
priority_queue<int> q;
int len = arr.size();
int i = 0;
for (; i < k && i < len; ++i) {
q.push(arr[i]);
}
for (; i < len; ++i) {
if (q.top() > arr[i]) {
q.pop();
q.push(arr[i]);
}
}
while (!q.empty()) {
ans.push_back(q.top());
q.pop();
}
return ans;
}
方法二:快速排序思想
可以借鉴快速排序思想中的划分函数。划分函数可以将一个数组分为两个部分,小于等于分界值 pivot 的元素都会被放到数组的左边,大于的都会被放到数组的右边,然后返回分界值的下标。
与快速排序不同的是,快速排序会根据分界值的下标递归处理两侧,而我们只需要处理一侧。
此时可以划分为三种情况:
- 如果分界值的下标等于 k - 1 或 k,则说明数组的第 k 小数是 pivot 或第 k+1 小数是 pivot,此时可以确定,前 k 个数就是数组的前 k 小值,将其保存到新数组中返回即可。
- 如果分界值的下标小于 k - 1,则说明数组的第 k 小数在分界值的右侧,此时需要递归划分分界值右侧的数组。
- 如果分界值的下标大于 k,则说明数组的第 k 小数在分界值的左侧,此时需要递归划分分界值左侧的数组。
实现中定义 partition(arr, l, r, k) 递归函数,递归完成后取数组的前 k 个数放到新数组中返回即可。
vector<int> smallestK_V2(vector<int>& arr, int k) {
if (arr.empty())
return {};
int l = 0;
int r = arr.size() - 1;
partition(arr, l, r, k);
vector<int> ans;
for (int i = 0; i < k; ++i) {
ans.push_back(arr[i]);
}
return ans;
}
void partition(vector<int>& arr, int l, int r, int k) {
int i = l - 1;
int x = arr[r];
for (int j = l; j < r; ++j) {
if (arr[j] <= x) {
++i;
std::swap(arr[j], arr[i]);
}
}
std::swap(arr[i + 1], arr[r]);
if (i + 1 == k || i + 1 == k - 1)
return;
else if (i + 1 < k - 1)
partition(arr, i + 2, r, k);
else if (i + 1 > k)
partition(arr, l, i, k);
}