前言
用堆解决TopK问题需要读者掌握堆排序的操作,如果有读者不理解堆排序的话欢迎先看完读者的另一篇文章再来阅读此篇文章:堆排序过程详解
解决思路概括
TopK问题就是要找出所给数据中最接近所给条件的K个数据,这个所给条件可以是最大值,最小值,或者是最短字符等等,当然,对于所给的不同条件,需要对应调整堆的结构,如果处理的是整型数据,堆元素的类型就要是整型,换成其它类型就得改变元素的数据类型,笔者这里就讲整型数据TopK的处理。
如果我们要找最大的K个数据,有些老铁可能会想:就遍历K遍数据,像冒泡那样每次找到的最大值都放后边不就行了?当然,这也是一种方法,但是还有优化的空间,其实,遍历一遍就可以找出来!是不是很惊讶?
构建TopK堆
现在假设我们有1亿个数据,要求我们找出最大的5个,我们先假设前5个数据即为最大的五个数,取数据前五个进行建堆,后续遍历如果出现大于堆中最小值的数据,就进行替换,这里建堆的时间复杂度可以忽略不计(因为建立一个5个元素的堆,相较于遍历1亿个数据,可谓是九牛一毛了)那么是建大堆还是小堆呢?如果建大堆,那么堆顶元素就会是最大的那个,堆底元素在堆中的相对大小不确定,建小堆的话堆顶元素为最小的那个,同样不能确定堆底元素的相对大小;前面说了堆的意义就是保存最大的5个数据,而每次遍历都要考虑替换掉堆中的最小值,为了每次都能找到堆中的最小值,我们就要建什么堆呢?没错,是小堆,小堆的堆顶元素即为堆中的最小值,这样每次遍历时只要与堆顶元素比较,就能判断是否要更新TopK堆(即刚建的5个元素的堆)那么该怎么进行更新呢?
不同的建堆思路对应不同的更新方法,这里的建堆思路不是指所建堆的类型,而是所建堆的位置,由于堆的存储结构是顺序结构,故有些老铁可能选择直接将堆的堆顶指针指向代处理的数组的首元素,我们就叫它“原地建堆”吧,而有些老铁可能没有这种想法,而是选择先读取待处理的数据,然后进行建堆,我们就叫它“异地建堆”吧;两种方法都可行,虽然前者(原地建堆)是好一些,但没必要纠结于那一点点细小的差异,因为复杂度都是O(1)级别的,先想到那种就写那种吧,但笔者还是更推荐原地建堆的方法,因为写起来让人感觉很舒服;当然,如果有特殊情况,比如待处理数据的顺序是有意义的,不可轻易改变,就得考虑异地建堆来处理了。
以下是两种思路的简单实现(具体实现可见文章结尾):
原地建堆
#include <stdlib.h>
#include <assert.h>
// 堆的定义
typedef int HPDataType;
typedef struct Heap {
int size; // 堆的元素个数
int capacity; // 堆的容量,因为是顺序存储结构,故需要考虑容量问题
HPDataType* arr; // 堆顶元素指针
}Heap;
// 原地建堆演示
int main(){
int nums[7] = {1,2,3,4,5,6,7};
int k = 7;
// 建堆
Heap* ph = (Heap*)malloc(sizeof(Heap));
assert(ph);
// 堆的属性设置
ph->arr = nums; // 这里是原地建堆
ph->capacity = ph->size = k;
// 建堆
int i = 0;
for (i = (k - 1 - 1) / 2; i >= 0; i--) { // 这里的k是位序,转化成数组下标时需要减一
AdjustDown(ph, k, i);
// 此处笔者建的是小堆,故此函数求得的是最大的K个数据,如要求最小数据,将向下调整改成建大堆模式即可
}
return 0;
}
异地建堆
#include <stdlib.h>
#include <assert.h>
// 堆的定义
typedef int HPDataType;
typedef struct Heap {
int size; // 堆的元素个数
int capacity; // 堆的容量,因为是顺序存储结构,故需要考虑容量问题
HPDataType* arr; // 堆顶元素指针
}Heap;
// 异地建堆的演示
int main(){
int nums[7] = {1,2,3,4,5,6,7};
int k = 3;
// 堆的创建
Heap* ph = (Heap*)malloc(sizeof(Heap));
assert(ph);
// 堆的属性设置
ph->arr = (int*)malloc(sizeof(int) * k);
ph->size = ph->capacity = k;
// 读取数据
int i = 0;
for (i = 0; i<k; i++){
ph->arr[i] = nums[i];
}
// 建堆
for (i = (k - 1 - 1) / 2; i >= 0; i--) { // 这里的k是位序,转化成数组下标时需要减一
AdjustDown(ph, k, i);
// 此处笔者建的是小堆,故此函数求得的是最大的K个数据,如要求最小数据,将向下调整改成建大堆模式即可
}
return 0;
}
遍历更新TopK堆
在建好TopK堆后,我们从待处理数据的第K+1个元素开始进行遍历,每次遍历都与堆顶元素进行比较,如果满足入堆条件,就进行一次更新,如此遍历一遍后,堆中就为最接近条件的K个元素了,然后我们说说更新思路吧:
原地建堆:由于使用了待处理数据的空间,为了待处理数据的完整性,更新TopK堆时就不能直接使用覆盖操作,而是要使用交换操作,这样才能保证待处理数据的完整性,然后再进行一次向下调整即可。
异地建堆:异地的空间就不用考虑堆元素被覆盖时,会对待处理数据产生影响的问题,我们直接用满足条件的遍历元素覆盖堆顶元素,然后进行一次向下调整即可。
代码汇总
用堆解决TopK问题需要堆的定义、建堆操作和堆的向下调整操作即可,具体代码如下:
#include <stdlib.h>
#include <assert.h>
typedef int HPDataType;
typedef struct Heap {
int size; // 堆的元素个数
int capacity; // 堆的容量,因为是顺序存储结构,故需要考虑容量问题
HPDataType* arr; // 堆顶元素指针
}Heap;
// 原地建堆版
void TopK(HPDataType* nums, int len, size_t k) {
assert(len >= k); // 这里等于k也进去是为了保证在进行一次TopK的调用后能形成k个元素的堆
Heap* ph = (Heap*)malloc(sizeof(Heap));
assert(ph);
ph->arr = nums; // 这里是原地建堆
ph->capacity = ph->size = k;
int i = 0;
for (i = (k - 1 - 1) / 2; i >= 0; i--) { // 这里的k是位序,转化成数组下标时需要减一
AdjustDown(ph, k, i);
// 此处笔者建的是小堆,故此函数求得的是最大的K个数据,如要求最小数据,将向下调整改成建大堆模式即可
}
for (i=k; i<len; i++) {
if (nums[i] > ph->arr[0]) {
// 因为是原地建堆,故要保障数据的完整性
// 如果是异地建堆,就“不能”进行交换!而是直接用遍历元素覆盖堆顶元素
// 因为交换会造成原数据被覆盖的现象,读者可以思考一下
// 此外异地建堆还要处理返回值的问题,即要将建好的TopK堆返回
Swap(&nums[i], &ph->arr[0]);
AdjustDown(ph, k, 0);
}
}
}
// 向下调整建小堆
void AdjustDown(Heap* ph, int n, int parent) {
// 参数解释:ph 是指向堆顶元素的指针,n代表此次向下调整的范围,相当于向下调整的终点,parent表示向下调整的起点
assert(ph);
int child = parent * 2 + 1;
if (child + 1 < n && ph->arr[child] > ph->arr[child + 1]) { // 如要建大堆,此处&&右边的>改成<,后续还有要修改的
child += 1;
}
while (child < n) {
if (ph->arr[child] < ph->arr[parent]) { // 如要建大堆,此处的>改成<,后续还有要修改的
Swap(&ph->arr[child], &ph->arr[parent]);
}
parent = child;
child = parent*2 + 1;
if (child + 1 < n && ph->arr[child] > ph->arr[child + 1]) { //如要建大堆,此处&&右边的>改成<,建大堆修改到这里完毕
child += 1;
}
}
}
// 交换堆元素
void Swap(HPDataType* p1, HPDataType* p2) {
HPDataType tmp = *p1;
*p1 = *p2;
*p2 = tmp;
}
///
// 异地建堆版
HPDataType* TopK(HPDataType* nums, int len, size_t k) {
assert(len >= k); // 这里等于k也进去是为了保证在进行一次TopK的调用后能形成k个元素的堆
Heap* ph = (Heap*)malloc(sizeof(Heap));
assert(ph);
// 堆的属性设置
ph->arr = (int*)malloc(sizeof(int) * k);
ph->size = ph->capacity = k;
// 读取数据
int i = 0;
for (i = 0; i<k; i++){
ph->arr[i] = nums[i];
}
// 建堆
for (i = (k - 1 - 1) / 2; i >= 0; i--) { // 这里的k是位序,转化成数组下标时需要减一
AdjustDown(ph, k, i);
// 此处笔者建的是小堆,故此函数求得的是最大的K个数据,如要求最小数据,将向下调整改成建大堆模式即可
}
for (i=k; i<len; i++) {
if (nums[i] > ph->arr[0]) {
// 异地建堆不用考虑覆盖的影响,直接覆盖即可
ph->arr[0] = nums[i];
AdjustDown(ph, k, 0);
}
}
return ph->arr; // 如果不接收这里的返回值,则会造成内存泄漏的问题
}
// 向下调整建小堆
void AdjustDown(Heap* ph, int n, int parent) {
// 参数解释:ph 是指向堆顶元素的指针,n代表此次向下调整的范围,相当于向下调整的终点,parent表示向下调整的起点
assert(ph);
int child = parent * 2 + 1;
if (child + 1 < n && ph->arr[child] > ph->arr[child + 1]) { // 如要建大堆,此处&&右边的>改成<,后续还有要修改的
child += 1;
}
while (child < n) {
if (ph->arr[child] < ph->arr[parent]) { // 如要建大堆,此处的>改成<,后续还有要修改的
Swap(&ph->arr[child], &ph->arr[parent]);
}
parent = child;
child = parent*2 + 1;
if (child + 1 < n && ph->arr[child] > ph->arr[child + 1]) { //如要建大堆,此处&&右边的>改成<,建大堆修改到这里完毕
child += 1;
}
}
}