TopK算法

TopK的应用

TopK主要解决的是从一组数据中选择出前K个最大或最小的元素的问题。在计算机科学和数据分析领域，TopK方法被广泛应用于排序、搜索和排名等任务，它对于处理海量数据特别有效。

TopK相较于堆排序的优势

也许你会觉得，我们通过前面的堆排序取K次堆顶元素也可以解决这种问题。但是当数据规模巨大或存在特定约束时，这种方法并不是十分适合。为什么呢？我们通过一个例子来理解：

例如，假设我们有一个包含100亿个整数的无序数组，需要找出其中最大的100个数。在这种情况下，如果直接使用堆排序，我们需要先构建一个大顶堆，然后依次取出堆顶元素（即当前最大值），直到取出100个元素为止。这个过程的时间复杂度为O(N log K)，其中N是数组的长度，K是需要找出的最大元素的数量。虽然这个时间复杂度在理论上是可接受的，但在实际应用中，由于N非常大（100亿个整数的无序数组会占用37.25G的内存），构建和维护这个大顶堆将消耗大量的内存和时间资源。

TopK解决问题的思路

假设我们要在一个包含n个整形的无序数组里找前K个最大值：

1. 将数组内前K个数据最小的元素取出建小顶堆。
2. 遍历数组中剩余的元素与小顶堆顶比。
3. 若当前元素大于当前堆顶的元素的数据则替换堆顶，再调堆。
4. 遍历完数组中剩余的数后，堆中即为TopK大值。

这里补充一句：找前K个最小的值需要将数组内前K个数据最大的元素取出建大堆。

代码

void swap(int* x, int* y)
{
	int tmp = *x;
	*x = *y;
	*y = tmp;
}

int getMin(int* a, int x, int y)
{
	if (a[x] < a[y])
	{
		return x;
	}
	else {
		return y;
	}
}

void adjustDown(int* a, int parent, int size)
{
	assert(a);

	while (parent < size)
	{
		int lChild = 2 * parent + 1, rChild = 2 * parent + 2, swapChild;
		if (rChild < size) {
			swapChild = getMin(a, lChild, rChild);
		}
		else {
			swapChild = lChild;
		}

		if (swapChild < size && a[parent] > a[swapChild]) {
			swap(&a[parent], &a[swapChild]);
			parent = swapChild;
		}
		else {
			break;
		}
	}
}

void printTopK(int* a, int n, int k) {
	//1.用a中前k个元素建小堆
	for (int i = (n - 2) / 2; i >= 0; i--) {
		adjustDown(a, i, n);
	}
	for (int i = 0; i < k; i++) {
		printf("%d ", a[i]);
	}
	printf("\n");
	//2.将剩余n-k个元素与堆顶元素比较大小，大于堆顶元素则替换
	int cnt = n - k, cur = k;
	while (cnt) {
		if (a[0] < a[cur]) {
			swap(&a[0], &a[cur]);
			adjustDown(a, 0, k);
		}
		--cnt, ++cur;
	}
	for (int i = 0; i < k; i++) {
		printf("%d ", a[i]);
	}
	printf("\n");
}

//测试案例
void testTopK() {
	int n = 100000, k = 10;
	int* a = (int*)malloc(sizeof(int) * n);
	if (a == NULL) {
		perror("malloc fail");
		return;
	}
	srand(time(0));
	for (int i = 0; i < n; i++) {
		a[i] = rand() % 10000;
	}
    //提前准备好最大的10个数，方便后面测试
	a[1] = 111111;
	a[6] = 1001177;
	a[27] = 20001;
	a[99] = 1000001;
	a[278] = 99999;
	a[678] = 23456;
	a[1001] = 777777;
	a[345] = 1100011;
	a[4567] = 30000;
	a[9999] = 712903;
	printTopK(a, n, k);
}