C++代码实现Top-K问题最优解决办法

Top-K问题

1、问题描述

Top-K问题是一个十分经典的问题，一般有以下两种方式来描述问题：在10亿的数字里，找出其中最大的100个数；或者在一个包含n个整数的数组中，找出最大的100个数。

前边两种问题描述稍有区别，但都是说的Top-K问题，前一种描述方式是说这里也许没有足够的空间存储大量的数字或其他东西，我们最好可以在一边输入数据，一边求出结果，而不需要存储数据；后一种说法则表示可以存储数据，这种情况下，最简单直观的想法就是对数组进行排序，取后100个数即为所求。

2、解法思想和实现

第一种情况

①首先说第一种情况的思路，这种情况下，关键在于不能消耗太大的内存，无法通过数组的简单排序来求取最大的K个数，于是我们应该想到堆排序，求最大的K个数，就采用大小为K的最小二叉堆来实现；我们知道二叉堆可以看作是一颗近似的完全二叉树，其根节点正好就是K个数中最小的一个。

具体算法：先输入K个数，建立一个大小为K的最小二叉堆，接着每输入一个数，与堆的根节点进行比较，如果比根节点还小，说明不可能为最大的K个数之一，如果比根节点大，那么替换根节点的值，接着下沉根节点，维护二叉堆的性质。这样到成功输入所有数据后，最小二叉堆里剩下的就为最大的K个数。（如果求最小的K个数，同理换成最大二叉堆即可）。

时间复杂度：由于算法主要涉及对二叉堆结构的操作，所以总体时间复杂度为O(nlgK)。

/*
*代码采用STL中的最小优先队列实现，由于STL中自带最小优先队列，其底层就是二叉堆实现，
*所以就不再手写二叉堆了。最小优先队列顶层元素总是队列中最小的元素，也就是二叉堆堆顶。
*/
 
#include <iostream>
#include <queue>
#include <vector>
using namespace std;
 
/*由于STL自带优先队列是默认最大优先的，所以自己写了一个比较函数，将其改为最小优先*/
struct cmp1 {
	bool operator ()(int &a, int &b) {
		return a>b;											//最小值优先
	}
};
 
int main() {
	//这里用来测试，输入格式：先输入需要求的最大K个数中的K值，再依次输入数据流
	int K = 0;
	cin >> K;
	int tmp = 0;
	int i = 0;
	priority_queue<int,vector<int>,cmp1> minHeap;			//建立最小优先队列
	while (cin >> tmp) {									//循环输入数据流
		if (i < K) {										//先建立一个K个大小的优先队列，也就是K大小的二叉堆
			minHeap.push(tmp);
		}
		else {												//算法实现
			if (tmp <= minHeap.top())
				continue;
			else if (tmp > minHeap.top()) {
				minHeap.pop();
				minHeap.push(tmp);
			}
		}
		i++;
	}
	while (!minHeap.empty()) {								//输出最大的K个数
		cout << minHeap.top() << endl;
		minHeap.pop();
	}
	return 0;
}
 

第二种情况

②第二种情况，这种情况，由于可以操作存储数据的数组，所以我们采用排序的方式进行求解，但一般的排序时间复杂度也挺高，题目只求最大的K个数，不需要完全排序；于是我们想到可以借用快排思想来进行求解。

这个解法源于快排（Quick Sort），所以也叫Quick Select，主要基于快排中Partition函数。

具体算法：我们知道，每运行一次Partition函数都会确定一个数m的最终位置，且小于m的数均在其左边，大于m的数都在其右边，所以我们的目的就是当数m的右边正好有K-1个数的时候停止算法，得到答案。每次运行Partition函数时，根据前边上述性质，若

• K<右边数组长度，那么要找的K个数一定在m右边，对m右边的数组运行Partition函数；

• K=右边数组长度+1，那么正好找到最大的K个数，为数m以及其右边数组，停止算法；

• 其他情况，最大的K个数不仅仅在m数右边数组中，对m左边数组运行Partition函数。

  时间复杂度：与快排类似，Quick Select同样也是不稳定的算法，最坏情况下时间复杂度会达到O(n^2)，但平均性能也与快排类似，时间复杂度一般可认为为O(n)。

#include <iostream>
#include <vector>

using namespace std;

void swap(vector<int>&arr, int l, int r)//元素交换函数
{
	int temp = arr[l];
	arr[l] = arr[r];
	arr[r] = temp;
}

int Partition(vector<int>&arr, int left, int right)
{
	int less = left - 1;//选准左边界
	int more = right;//右边界
	int temp = arr[right];//选定基准位置
	while (left < more)
	{
		if (arr[left] < temp)
		{
			swap(arr, ++less, left++);//当前元素小于基准元素，左边界内扩
		}
		else if (arr[left] > temp)//当前元素大于基准元素，右边界内扩，左边界不变
		{
			swap(arr, --more, left);
		}
		else
			left++;
	}
	swap(arr, more, right);//最后一个基准位置交换
	return left;//如果存在元素相等情况，返回相同元素两侧的边界索引
}


int main() {
	cout << "请输入K： " << endl;
	int K = 0;										//测试部分，输入需要求的K值大小，然后再依次输入数组元素
	cin >> K;
	int tmp = 0;
	vector<int> vec = {1,2,6,8,10,50,34,36,27,58,70,66};
	
	int size = vec.size();
	if (size == 0 || K > size)
		return 0;
	if (size == K)
	{
		for (int i = 0; i < size; i++) {			//测试部分，输出需要求的K个数
			cout << vec[i] << endl;
		}
	}
	int left = 0;
	int right = vec.size() - 1;
	int index = Partition(vec, left, right);
	while (index != size - K) {						//当Partition返回值及右边部分不是K大小时，继续循环
		int sizeOfRight = size - index - 1;			//记录index右边数组长度大小
		if (K <= sizeOfRight) {
			index = Partition(vec, index + 1, right);
		}
		else if (K == sizeOfRight + 1)				//这一步好像有点多余，while循环保证了这点，但为了对应博客文字描述就加上了
			continue;
		else if (K > sizeOfRight + 1) {
			index = Partition(vec, left, index - 1);
		}
	}
	cout << "输出TOPK个元素： " << endl;
	for (int i = index; i < size; i++) {			//测试部分，输出需要求的K个数
		cout << vec[i] << " ";
	}
	cout << endl;
	return 0;
}

3、 快速排序找第K大的数

#include <iostream>
#include <vector>

using namespace std;

void swap(vector<int>&arr, int l, int r)//元素交换函数
{
	int temp = arr[l];
	arr[l] = arr[r];
	arr[r] = temp;
}

int Partition(vector<int>&arr, int left, int right)
{
	int less = left - 1;//选准左边界
	int more = right;//右边界
	int temp = arr[right];//选定基准位置
	while (left < more)
	{
		if (arr[left] < temp)
		{
			swap(arr, ++less, left++);//当前元素小于基准元素，左边界内扩
		}
		else if (arr[left] > temp)//当前元素大于基准元素，右边界内扩，左边界不变
		{
			swap(arr, --more, left);
		}
		else
			left++;
	}
	swap(arr, more, right);//最后一个基准位置交换
	return left;//如果存在元素相等情况，返回相同元素两侧的边界索引
}


int main() {
	cout << "请输入K： " << endl;
	int K = 0;										//测试部分，输入需要求的K值大小，然后再依次输入数组元素
	cin >> K;
	int tmp = 0;
	vector<int> vec = { 1,2,6,8,10,50,34,36,27,58,70,66 };

	int size = vec.size();
	if (size == 0 || K > size)
		return 0;
	if (size == K)
	{
				//测试部分，输出需要求的K个数
			cout << vec[K - 1] << endl;

	}
	int left = 0;
	int right = vec.size() - 1;
	int index = Partition(vec, left, right);
	while (index != K - 1) {						//当Partition返回值及右边部分不是K大小时，继续循环
		//int sizeOfRight = size - index - 1;			//记录index右边数组长度大小
		if (  index < K - 1) {
			index = Partition(vec, index + 1, right);
		}
		else if (index > K - 1) {
			index = Partition(vec, left, index - 1);
		}
	}

	cout << vec[index] << endl;
	return 0;
}