【堆、快速选择排序】探寻TopK问题的解决方案

目录

前言

什么是TopK问题

建堆——优先级队列

快速选择排序QuickSelect

快速选择排序的时间复杂度

---

前言

TopK问题在面试中常常被问到 —— 比如，在10亿个整数里，找出最大的前100个。在海量数据中查找出重复出现的元素或者去除重复出现的元素也是常考的问题。所以在对待TopK问题上，我们可不敢掉以轻心。下面，我们进入正题。

什么是TopK问题

TopK问题就是在一个数据集合中找出最大的前K个或者最小的前K个。在面试中遇到这个问题时，通常它的数据量很大，动不动就上亿或者海量数据。这就导致我们无法直接在内存中对所有数据进行排序或者说排序的时间成本很高等等，因此这类问题通常不采取直接排序再查找的做法。下面我们看看几种常见的做法。

建堆——优先级队列

以这道题为例——在10亿个整数里，找出最大的前100个。下面我们看看堆是如何解决这个问题的。

用100个数建一个小堆，然后将剩余数据依次与堆顶元素比较，比堆顶数据大的则替换堆顶数据进堆，遍历完剩余的数据后，堆里面的值就是最大的前100个。会有同学在这里可能有这样的疑问——第50大的数据在堆顶时会不会挡住第51大、第52大的数据进堆。当然是不会的。如果第50大的数据经过向下调整后在堆顶，那么说明堆里面剩下的99个数都要比第50大的数大。仔细体会一下，这是不是很荒谬——第50大数据：我都已经是第50大了，居然还有99个数比我还大，看不起谁呢？！所以说前100大的数据是不会被挡住的。

我们先来重温优先级队列priority queue，下面是摘自文档的关于priority queue的解释：

Priority queues are a type of container adaptors, specifically designed such that its first element is always the greatest of the elements it contains, according to some strict weak ordering criterion.

可以看出，priority queue默认情况下是大堆，而我们这里需要一个小堆，故需要我们自己写一个比较方法。下面我们开始写代码：

#include <iostream>
#include <queue>
using namespace std;

#define N 10000  //数据个数
#define K 10	 //要找的前K个

//比较方法——因为默认是大堆，而我们的需求是小堆，故需传比较方法
struct MyGreater
{
	bool operator()(const int& left, const int& right)
	{
		return left > right;
	}
};

//造数据
void CreateData()
{
	FILE* fin = fopen("data.txt", "w");
	for (size_t i = 0; i < N; i++)
	{
		//这里模上10000，那么文件中的数据都小于等于10000，
		//当数据生成完毕后，我随机选取了10个，在后边加了个0
		//这些数据就是最大的前10个，如果我们的程序运行结果
		//为这10个数，就说明我们程序是对的，其实就是为了
		//方便确认程序的结果得到的是否为最大的前10个
		fprintf(fin, "%d\n", (rand() + i) % 10000);
	}
	fclose(fin);
}

//从文件中读取数据
void PrintData()
{
	priority_queue<int, vector<int>, MyGreater> pq;
	FILE* fout = fopen("data.txt", "r");
	int num = 0;

	//取文件中的前K个数据来进小堆
	for (size_t i = 0; i < K; i++)
	{
		fscanf(fout, "%d", &num);
		pq.push(num);
	}

	//读取剩下的 N - K 个数据
	while (!feof(fout))
	{
		fscanf(fout, "%d", &num);
		if (num > pq.top())
		{
			pq.pop();
			pq.push(num);
		}
	}
	fclose(fout);

	while (!pq.empty())
	{
		int top = pq.top();
		pq.pop();
		cout << top << " ";
	}
	cout << endl;
}

int main()
{
	srand(time(NULL));
	//CreateData();
	PrintData();
	return 0;
}

运行结果： 

注意，在造数据时，我们应当先注释掉PrintData()，当造完数据后，打开数据所在文件进行改造，再注释掉CreateData()，调用PrintData()，避免重复造数据，导致我们所对数据的改造失效。还有一点，我这里只是在10000个数据中找到最大的前10个，如果你有兴趣可以把数据量在调大一些，整个上亿个数据也不是不行，只不过运行时间可能长一些罢了。 

快速选择排序QuickSelect

快速选择排序和快速排序可以说是孪生兄弟了，这两个算法均出自hoare大佬之手，它们的思想尤为相似。快排一上来先确定一个基准值key，基准值key的左边均小于等于key，右边均大于key。接着再分别对左右子数组进行同样的操作，直到左右子数组不能再分割（即只有一个元素）。快速选择排序一上来也是先确定一个基准值key，然后进行一趟快排，使基准值key的左边均小于等于key，右边大于key，然后再对目标子数组进行分割，直到找到目标值。快排与快速选择排序的区别就在于后者不在对左右子数组都进行处理，而是针对目标子数组进行处理(可能是左子数组也可能是右子数组)。温故知新，既然快速选择排序和快排强相关，我们不妨把它俩一块写了。

//三数取中
int GetMidIndex(int* arr, int left, int right)
{
	int mid = (left + right) / 2;
	if (arr[left] > arr[right])
	{
		if (arr[right] > arr[mid])
			return right;
		else if (arr[mid] < arr[left])
			return mid;
		else
			return left;
	}
	else 
	{
		if (arr[left] > arr[mid])
			return left;
		else if (arr[mid] < arr[right])
			return mid;
		else
			return right;
	}
}

//对数组进行分割
int partition(int* arr, int left, int right)
{
	int index = GetMidIndex(arr, left, right);
	swap(arr[left], arr[index]);
	int keyi = left;
	while (left < right)
	{
		while (left < right && arr[right] >= arr[keyi]) --right;
		while (left < right && arr[left] <= arr[keyi]) ++left;
		swap(arr[left], arr[right]);
	}
	swap(arr[keyi], arr[left]);
	return left;
}

//快排
void quickSort(int* arr, int left, int right)
{
	if (left >= right)
		return;
	int mid = partition(arr, left, right);
	quickSort(arr, 0, mid - 1);
	quickSort(arr, mid + 1, right);
}

//快速选择排序
//返回值为第k大的数据，例如k=1,则返回最大的数据
int quickSelect(int* arr, int left, int right, int k)
{
	if (left >= right)
		return arr[right];
	int index = partition(arr, left, right);//index为基准值的下标
	if (right - index + 1 > k)
		return quickSelect(arr, index + 1, right, k);
	else if (right - index + 1 == k)
		return arr[index];
	else
		return quickSelect(arr, left, index - 1, k - right + index - 1);
}

int main()
{
	int arr[] = { 10,3,6,1,5,0,9,2,4,8,7 };
	size_t n = sizeof(arr) / sizeof(int);
	quickSort(arr, 0, n - 1);

	cout << "Whole sequence: ";
	for (size_t i = 0; i < n; i++)
	{
		cout << arr[i] << " ";
	}
	cout << endl;

	//我们可以找出单个，也可找出多个
	cout << "Second largest:" << quickSelect(arr, 0, n - 1, 2) << endl;

	//前3大
	cout << "Top 3: ";
	for (size_t i = 1; i <= 3; i++)
	{
		cout << quickSelect(arr, 0, n - 1, i) << " ";
	}
	cout << endl;

	return 0;
}

运行结果：

可以看到，快排和快速选择排序完全复用一个数组分割函数。快速选择排序算法既可以只拿到一个数据，例如找出第2大的数，也可以找出前k大的数，例如上面代码找出前3大的数据。也许你看到了这里，对上面快速选择排序的代码会有些迷糊，这不是你的错，是我的错，因为顺手写了快排，为了让它们在一起，就把快速选择排序的代码给出了。下面，我来解释解释这段代码。

首先明确一点，上面我所使用的分割数组函数是：基准值左边的值均小于等于基准值，基准值右边均大于基准值（其实等于基准值的值在左边还是右边都无所谓）。

以红色为基准，紫色为左子数组，粉色为右子数组。 假设要在上图中找第10大的数据，且此时左子数组的值均小于等于基准值，右子数组的值均大于基准值。

下标3到下标11总共有9个数据（这九个数据都大于等于左子数组的值），说明第10大的数据必然在左子数组中，且应为左子数组中的第一大，因为前9大为下标3到下标11的9个数据。

故当 right - index + 1 < k 时，第k大的数必然在左子数组中，且为左子数组中的第 k - (right - index + 1) 大。

假设现在要在上图中找出第9大数据，且此时左子数组的值均小于等于基准值，右子数组的值均大于基准值。从index到right正好有9个数据，而且这9个是数就是前9大数据，index右边有8个比它大的数，所以下标为index的元素就是第9大数据。换句话说，就是index到right为前9大，下标为4到下标为11的为前8大，那么下标为3的必然就是第9大。

故当 right - index + 1 == k 时，第k大数据为arr[index]。

假设现在要在上图中找出第5大数据，且此时左子数组的值均小于等于基准值，右子数组的值均大于基准值。从下标4到下标11有8个数据，而且毫无疑问这8个数据就是前8大数据，所以第5大的数据必然在这8个数当中，而且还是这8个数当中的第5大。

故当 right - index + 1 > k 时，第k大数据必然出现在右子数组中，且仍为第k大。

快速选择排序的时间复杂度

快速选择排序的时间复杂度推导起来比较困难，有兴趣的话可以翻阅《算法导论》，这里就直接给出结论：O(n)。

---

本文到这就结束啦，感谢你的支持！