在刷题或者面试的过程中,我们经常会遇到这样一种题目,给一个数组,求最大或者最小的 k 个数,或者第 k 大或者第 k 小的数字
解法看到这样一个题目,首先想到的就是大小堆,创建大小为 k 的大根堆或者小根堆,然后不断的更新堆,下面以求最大的 k 个数字为例,具体代码如下
void BuildHeap(std::vector<int>& arr, int index, int len) {
int i = index * 2 + 1;
while (i < len) {
if ((i + 1) < len && arr[i + 1] < arr[i]) {
i = i + 1;
}
if (arr[i] > arr[index]) {
break;
}
std::swap(arr[i], arr[index]);
index = i;
i = 2 * index + 1;
}
}
void MakeHeap(std::vector<int>& arr) {
int mid = arr.size() / 2;
for (int i = mid; i >= 0; --i) {
BuildHeap(arr, i, arr.size());
}
}
std::vector<int> MaxK(int* arr, int len, int k) {
std::vector<int> vec;
for (int i = 0; i < k; ++i) {
vec.push_back(arr[i]);
}
MakeHeap(vec);
for (int i = k; i < len; ++i) {
if (arr[i] > vec[0]) {
vec[0] = arr[i];
BuildHeap(vec, 0, k);
}
}
return vec;
}
int main(void) {
int arr[] = {1,8,3,4,5};
std::vector<int> val = MaxK(arr, sizeof(arr) / sizeof(arr[0]), 3);
for (int i = 0; i < 3; ++i) {
std::cout << val[i] << " ";
}
std::cout << std::endl;
getchar();
return 0;
}
另一种解法 std::priority_queue
C++ STL 有一个优先级队列 std::priority_queue,它的内部就是使用堆来实现的,接下来我们使用 std::priority_queue 在来一次
std::vector<int> MaxK(int* arr, int len, int k) {
std::priority_queue<int, std::vector<int> , std::greater<int>> q{arr, arr+k};
for (int i = k; i < len; ++i) {
if (arr[i] > q.top()) {
q.push(arr[i]);
q.pop();
}
}
std::vector<int> result;
while(q.size() > 0) {
result.push_back(q.top());
q.pop();
}
return result;
}
int main(void) {
int arr[] = {1,8,3,4,5};
std::vector<int> val = MaxK(arr, sizeof(arr) / sizeof(arr[0]), 3);
for (int i = 0; i < val.size(); ++i) {
std::cout << val[i] << " ";
}
std::cout << std::endl;
getchar();
return 0;
}
代码瞬间精简了很多,可见优先级队列还是很好使用的
怎么创建优先级队列1、默认的创建
std::priority_queue<int> q{ arr, arr + k };
上面代码表示创建了一个优先级队列,初始化元素为数组的前 k 个元素,然后里面的元素从大到小排列,如果想要计算最小的 k 个数,或者第 k 小的数组,就需要按照这种方式创建堆
2、降序排列的优先级队列
std::priority_queue<int, std::vector<int> , std::greater<int>> q{arr, arr+k};
表示我们创建一个优先级队列,初始化元素为数组的前 k 个元素,排序的方式为从小到大排序,适用于求前 k 大的数字,或者第 k 大的数字