数据结构与算法：大小根堆和快速排序 解决TopK问题

问题：求出一组序列中值最小的前K个元素

在正常做法中，是先对这组元素排序，然后得出前k个元素，但排序算法的时间复杂度为O(n^2)或O(nlogn)，问题就是 能不能在线性时间内找到top k的元素呢？

文章目录

• 大小根堆
• • 求最小前k个元素
  • 代码示例
  • 求最大的前k个元素
  • 代码示例
  • 统计重复次数最少的k个元素
• 快排分割求TopK

大小根堆 求最小前k个元素

给出一组元素，求最小的前三个元素。

需要用到一个大根堆，堆顶元素是最大的。
思想：把大根堆堆顶的大值不断淘汰，放入小值。

图示：

这里来看为什么统计最小值要用大根堆，首先拿前k（这里取3）个元素组成大根堆，然后继续向后遍历数组，11比堆顶小，那么堆顶元素出堆，将11添加进来，并重新调整为大顶堆；再往后遍历，6比堆顶小，堆顶元素出堆，将6添加进来，再次调整为大顶堆…
这个示例的后续元素都比堆顶元素大，所以遍历时不做处理了，最终结果已经得出。

但是有人会问，这个调整为大根堆，也会耗时啊，但实际上这个堆的元素只有k个，每次调整时，时间复杂度最坏为 O(logK)，是一个常数级别的时间复杂度，再乘数组遍历的时间复杂度O(n)，那么时间复杂度为 O(n * logK)，logk为常数，可以省略，最终时间复杂度为 O(n)。

代码示例

#include 
#include 
#include 
#include 

using namespace std;

int main(void)
{
	vector<int> nums;
	srand(time(NULL));
	for (int i = 0; i < 1000; ++i)
	{
		nums.push_back(rand() % 10000 + 1);
	}

	// 求nums中值最小的前5个元素
	priority_queue<int> maxheap;
	int k = 5;

	// 前k个构建一个大根堆
	for (int i = 0; i < k; ++i)
	{
		maxheap.push(nums[i]);
	}

	// 遍历剩余元素
	for (int i = k; i < nums.size(); ++i)
	{
		if (maxheap.top() > nums[i])
		{
			maxheap.pop();
			maxheap.push(nums[i]);
		}
	}

	// 输出结果
	while (!maxheap.empty())
	{
		cout << maxheap.top() << " ";
		maxheap.pop();
	}
	cout << endl;

	return 0;
}

求最大的前k个元素

和上面相反，求最大元素，就应该使用小根堆，堆顶元素值最小。

思想：用k个元素构建小根堆，把堆顶小元素不断淘汰，把大元素放进来。

代码示例

#include 
#include 
#include 
#include 
using namespace std;

int main(void)
{
	vector<int> nums;
	srand(time(NULL));
	for (int i = 0; i < 1000; ++i)
	{
		nums.push_back(rand() % 10000 + 1);
	}

	// 求nums中值最大的前5个元素
	priority_queue<int> minheap;
	int k = 5;

	// 前k个构建一个大根堆
	for (int i = 0; i < k; ++i)
	{
		minheap.push(nums[i]);
	}

	// 遍历剩余元素
	for (int i = k; i < nums.size(); ++i)
	{
		if (minheap.top() < nums[i])
		{
			minheap.pop();
			minheap.push(nums[i]);
		}
	}
	// 输出结果
	while (!minheap.empty())
	{
		cout << minheap.top() << " ";
		minheap.pop();
	}
	cout << endl;

	return 0;
}

统计重复次数最少的k个元素

哈希表 + 大根堆

原理也很简单，先利用哈希表存储key-value，value为某元素出现次数，再利用大根堆，利用lambda表达式重写一个判断语句，最后遍历哈希表，找到结果。

#include 
#include 
#include 
#include 
#include 
#include 
using namespace std;

int main(void)
{
	vector<int> nums;
	srand(time(NULL));
	for (int i = 0; i < 1000; ++i)
	{
		nums.push_back(rand() % 10000 + 1);
	}

	int k = 3;
	// 统计重复次数最小的前3个数字
	unordered_map<int, int> map;
	for (auto key : nums)
	{
		map[key]++;
	}

	// 放入大根堆时，需要放key-value键值对
	using Type = pair<int, int>;
	using Cmp = function<bool(Type&, Type&)>;
	priority_queue<Type, vector<Type>, Cmp> maxheap(
		[](Type& a, Type& b)->bool {
			return a.second < b.second;
		}
	);

	auto it = map.begin();
	for (int i = 0; i < k; ++i, ++it)
	{
		maxheap.push(*it);
	}

	for (; it != map.end(); ++it)
	{
		if (maxheap.top().second > it->second)
		{
			maxheap.pop();
			maxheap.push(*it);
		}
	}

	while (!maxheap.empty())
	{
		cout << "key: " << maxheap.top().first
			<< " cnt: " << maxheap.top().second << endl;
		maxheap.pop();
	}

	return 0;
}

那么，统计重复次数最大的k个元素，和上面思想一样，使用哈希表 + 小根堆。

#include 
#include 
#include 
#include 
#include 
#include 
using namespace std;

int main(void)
{
	vector<int> nums;
	srand(time(NULL));
	for (int i = 0; i < 1000; ++i)
	{
		nums.push_back(rand() % 1000 + 1);
	}

	// 统计重复次数最大的前3个数字
	int k = 3;
	unordered_map<int, int> map;
	for (auto key : nums)
	{
		map[key]++;
	}

	// 放入大根堆时，需要放key-value键值对
	using Type = pair<int, int>;
	using Cmp = function<bool(Type&, Type&)>;
	priority_queue<Type, vector<Type>, Cmp> minheap(
		[](Type& a, Type& b)->bool {
			return a.second > b.second;
		}
	);

	auto it = map.begin();
	for (int i = 0; i < k; ++i, ++it)
	{
		minheap.push(*it);
	}

	for (; it != map.end(); ++it)
	{
		if (minheap.top().second < it->second)
		{
			minheap.pop();
			minheap.push(*it);
		}
	}

	while (!minheap.empty())
	{
		cout << "key: " << minheap.top().first
			<< " cnt: " << minheap.top().second << endl;
		minheap.pop();
	}
	return 0;
}

快排分割求TopK

如果对快速排序了解较少，可以看这一篇文章：【点此查看】

原理就是利用快速排序的划分函数Partition，该函数每次分割完，在基准值左边的元素，都是比基准值小的，右边全是比基准值大的元素，利用基准值的位置和k进行比较，很方便得到TopK元素。

#include 
#include 

using namespace std;

// 分割函数
int Partition(int* nums, int left, int right)
{
	int key = nums[left];

	while (left < right)
	{
		while (left < right && nums[right] > key)
			right--;

		if (left < right)
		{
			nums[left] = nums[right];
			left++;
		}

		while (left < right && nums[left] < key)
			left++;
		
		if (left < right)
		{
			nums[right] = nums[left];
			right--;
		}
	}

	nums[left] = key;
	return left;
}

// 求Topk函数
void selectTopK(int* nums, int left, int right, int k)
{
	int pos = Partition(nums, left, right);
	if (pos == k - 1)
	{
		return;
	}
	else if (pos > k - 1)
	{
		selectTopK(nums, left, pos - 1, k);
	}
	else
	{
		selectTopK(nums, pos + 1, right, k);
	}
}

void Show(int* nums, int size)
{
	for (int i = 0; i < size; ++i)
	{
		cout << nums[i] << " ";
	}
	cout << endl;
}

int main(void)
{
	int nums[] = { 64, 45, 52, 80, 66, 68, 0, 2, 18, 75 };
	int size = sizeof(nums) / sizeof(nums[0]);
	Show(nums, size);

	// 求值最小的前三个元素
	int k = 3;
	selectTopK(nums, 0, size - 1, k);
	Show(nums, k);
	return 0;
}