初阶数据结构之堆的应用——Top-K

TOP-K问题：即求数据结合中前K个最大的元素或者最小的元素，一般情况下数据量都比较大。

对于Top-K问题，能想到的最简单直接的方式就是：

排序 --------时间复杂度：O(N*logN)空间复杂度O(1)

建立N个数的大堆，Pop K次，就可以选出最大的前K个--------时间复杂度O(N+logN*K)空间复杂度：O(1)

但是：如果数据量非常大，排序就不太可取了(可能数据都不能一下子全部加载到内存中)。

最佳的方式就是用堆来解决，基本思路如下：

• 用数据集合中前K个元素来建堆

前k个最大的元素，则建小堆

前k个最小的元素，则建大堆

原理在堆排序已说

• 用剩余的N-K个元素依次与堆顶元素来比较，不满足则替换堆顶元素

将剩余N-K个元素依次与堆顶元素比完之后，堆中剩余的K个元素就是所求的前K个最小或者最大的元素。

这里采用向下调整法建小堆，因为这既方便效率又高。

//向下调整法
void AdjustDown(int*a, int& n, int parent)
{
	int child = parent * 2 + 1;
	while (child < n)
	{
		if (a[child] > a[child + 1] && child + 1 < n)
		{
			child++;
		}
		
		if (a[child] < a[parent])
		{
			swap(a[child], a[parent]);
			parent = child;
			child = parent * 2 + 1;
		}
		else
		{
			break;
		}
	}
}

然后就是解下来建堆：

    int* heapmink = (int*)malloc(sizeof(int) * k);
	assert(heapmink);
    //将数据赋值进创建的堆的中
	for (int i = 0; i < k; i++)
	{
		heapmink[i] = a[i];
	}
    //将数据按小堆建好
	for (int i = (k-1-1)/2; i >= 0; i--)
	{
		AdjustDown(heapmink, k, i);
	}

然后进行元素的交换：

    for (int j = k; j < n; j++)
	{
		if (heapmink[0] < a[j])
		{
			swap(heapmink[0], a[j]);
			//交换完后需要调整，将数据大的下调
			AdjustDown(heapmink, k, 0);
		}
	}

这样就可以将最大（最小的）k个数据给找到

以下是整个top-k解决的代码，自己实现了数据并进行测试：

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cassert>

using namespace std;

//交换
void swap(int& left, int& right)
{
	int temp = left;
	left = right;
	right = temp;
}

//向下调整法
void AdjustDown(int*a, int& n, int parent)
{
	int child = parent * 2 + 1;
	while (child < n)
	{
		if (a[child] > a[child + 1] && child + 1 < n)
		{
			child++;
		}
		
		if (a[child] < a[parent])
		{
			swap(a[child], a[parent]);
			parent = child;
			child = parent * 2 + 1;
		}
		else
		{
			break;
		}
	}
}


void PrintTopK(int* a, int& n, int k)
{
	// 1. 建堆--用a中前k个元素建堆
	int* heapmink = (int*)malloc(sizeof(int) * k);
	assert(heapmink);

	for (int i = 0; i < k; i++)
	{
		heapmink[i] = a[i];
	}

	for (int i = (k-1-1)/2; i >= 0; i--)
	{
		AdjustDown(heapmink, k, i);
	}


	// 2. 将剩余n-k个元素依次与堆顶元素交换，不满则则替换
	for (int j = k; j < n; j++)
	{
		if (heapmink[0] < a[j])
		{
			swap(heapmink[0], a[j]);
			//交换完后需要调整，将数据大的下调
			AdjustDown(heapmink, k, 0);
		}
	}

	//打印最大的k个数据
	for (int i = 0; i < k; ++i)
	{
		cout << heapmink[i] << " ";
	}
	cout << endl;

	//释放开辟的空间
	free(heapmink);
}

void TestTopk()
{
	int n = 10000;
	int* a = (int*)malloc(sizeof(int) * n);
	assert(a);
	srand(time(0));
	for (size_t i = 0; i < n; ++i)
	{
		a[i] = rand() % 1000000;
	}
	a[5] = 1000000 + 1;
	a[1231] = 1000000 + 2;
	a[531] = 1000000 + 3;
	a[5121] = 1000000 + 4;
	a[115] = 1000000 + 5;
	a[2335] = 1000000 + 6;
	a[9999] = 1000000 + 7;
	a[76] = 1000000 + 8;
	a[423] = 1000000 + 9;
	a[3144] = 1000000 + 10;
	PrintTopK(a, n, 10);
}

int main()
{
	TestTopk();
	return 0;
}