解决topK问题是寻找给定数据集中前K个最大或最小的元素。
常见有三种算法:
- 堆排序
- 维护一个大小为K的最小(或最大)堆。 遍历数据集,将元素插入堆中,如果堆大小超过K,则删除堆顶元素。
- 遍历结束后,堆中剩余的K个元素就是前K个最小(或最大)的元素。 时间复杂度:O(NlogK),其中N为数据集大小。
示例代码如下:
#include <iostream>
#include <vector>
#include <queue>
std::vector<int> topKHeap(std::vector<int>& nums, int k) {
std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap;
for (int num : nums) {
minHeap.push(num);
if (minHeap.size() > k) {
minHeap.pop();
}
}
std::vector<int> topk;
while (!minHeap.empty()) {
topk.push_back(minHeap.top());
minHeap.pop();
}
return topk;
}
int main() {
std::vector<int> nums = { 3, 2, 5, 2, 3, 2, 5, 5, 6 };
int k = 3;
std::vector<int> result = topKHeap(nums, k);
std::cout << "前 " << k << " 个最大的数字为:";
for (int num : result) {
std::cout << num << " ";
}
std::cout << std::endl;
return 0;
}
输出结果
前 3 个最大的数字为:5 5 6
- 快速排序
- 选择一个pivot元素将数据集划分为两部分:小于pivot的元素和大于pivot的元素。
- 当pivot的位置恰好在K-1时,那么pivot左边的元素就是前K个最小(或右边的元素是前K个最大)的元素。
- 如果pivot的位置小于K-1,继续在pivot右边的部分进行快速选择。
- 如果pivot的位置大于K-1,继续在pivot左边的部分进行快速选择。
- 时间复杂度:平均情况下为O(N),最坏情况下为O(N^2)。
这里给出参考文章地址 添加链接描述
#include <iostream>
#include <vector>
#include <algorithm>
int partition(std::vector<int>& nums, int low, int high) {
int pivot = nums[high];
int i = low - 1;
for (int j = low; j < high; j++) {
if (nums[j] < pivot) {
i++;
std::swap(nums[i], nums[j]);
}
}
std::swap(nums[i + 1], nums[high]);
return i + 1;
}
void quickSelect(std::vector<int>& nums, int low, int high, int k) {
if (low < high) {
int pivotIndex = partition(nums, low, high);
if (pivotIndex == k - 1) {
return;
}
else if (pivotIndex > k - 1) {
quickSelect(nums, low, pivotIndex - 1, k);
}
else {
quickSelect(nums, pivotIndex + 1, high, k);
}
}
}
std::vector<int> topKQuickSelect(std::vector<int>& nums, int k) {
quickSelect(nums, 0, nums.size() - 1, k);
std::vector<int> topk(nums.begin(), nums.begin() + k);
return topk;
}
int main() {
std::vector<int> nums = { 3, 2, 5, 2, 3, 2, 5, 5, 6 };
int k = 3;
std::vector<int> result = topKQuickSelect(nums, k);
std::cout << "前 " << k << " 个最小的数字为:";
for (int num : result) {
std::cout << num << " ";
}
std::cout << std::endl;
return 0;
}
结果为:
前 3 个最小的数字为:2 2 2
- 计数排序
计数排序(Counting Sort)是一种线性时间复杂度的排序算法,适用于输入数据范围较小的情况。解决topk问题可以使用计数排序的思想,具体步骤如下:
统计每个元素出现的次数:创建一个计数数组count,长度为输入数据范围的上限+1,初始值都为0。遍历输入数据,将每个元素出现的次数存储在相应的计数数组位置上。
累加计数数组:对计数数组进行累加操作,即将每个位置的值加上前一个位置的值。这一步的目的是确定每个元素在排序后的数组中的位置。
创建结果数组:创建一个与输入数组相同长度的结果数组result。
遍历输入数组并排序:从后向前遍历输入数组,根据输入元素在计数数组中的值,确定元素在结果数组中的位置,并将元素放入结果数组中。
返回结果:返回结果数组中的前k个元素。
#include <iostream>
#include <vector>
#include <algorithm>
std::vector<int> topK(std::vector<int>& nums, int k) {
int maxNum = *max_element(nums.begin(), nums.end());
std::vector<int> count(maxNum + 1, 0);
// 统计每个数字的出现次数
for (int num : nums) {
count[num]++;
}
// 找到出现次数最多的前 k 个数字
std::vector<int> topk;
for (int i = 0; i < k; i++) {
auto maxIt = std::max_element(count.begin(), count.end());
int maxIndex = std::distance(count.begin(), maxIt);
topk.push_back(maxIndex);
*maxIt = 0; // 将当前最大值置为0,避免重复选择
}
return topk;
}
int main() {
std::vector<int> nums = { 3, 2, 5, 2, 3, 2, 5, 5, 6 };
int k = 3;
std::vector<int> result = topK(nums, k);
std::cout << "出现次数前 " << k << " 的数字为:";
for (int num : result) {
std::cout << num << " ";
}
std::cout << std::endl;
return 0;
}
出现次数前 3 的数字为:2 5 3