在实际工作中我们经常会遇到将一个list中最大[最小]的前TopK个元素输出的问题。
比如说在电商领域,求上个月卖的最好的前10个商品,或者是每个品类下卖的最好的前10个商品。
这类问题,很多数据库或数据仓库工具可以提供直接的实现,比如第一个问题在mysql中就直接order by + limit就好; 而第二个问题在Hive中也很简单, distribute + order by + rownum + where也可以直接搞定,在Oracle则用 partition 就好。
而这些实现的原理都是基于对原list做一次排序,然后在排序结果的基础上取出前TopK个元素。
例如Python中可以这样:
a=[2,1,3,8,7,6,4,7,8,0,-1,3,2,6,3,9]
a.sort()
top10 = a[0:10]
其中的排序过程使得最后取出的前TopK个元素不仅能满足要求,而且还是排好序的。这似乎是一个非常不错的addtional benifit,不过你除非你确信或可接受由此带来的性能开销,最好不要轻易的接受这种不速的benifit。因为天下没有免费的午餐。。当然,就我见到的应用项目中,99%的做法还真就是这种sort + top的实现。没什么不好,只是他们比较“富裕”,能负担的起这个附加的午餐费而已。
不过,当有一天,我也用上述方法来解决我项目中的问题时发现,程序运行的很慢很慢,因为我的数据量很大很大。。。初步估计要运行好几百个小时,这当然是不能忍的。当然也是因为Python这种解释性语言本身比较低效。
所以必须换一种语言,顺道也换一种更经济的实现策略,就是本文接下来要介绍的:基于堆的topK实现。
很多同学可能都比较熟悉堆排序,但我们的问题并不需要排序,一旦排序,不管用什么方式,都会重蹈覆辙。
让我们重新回顾一下问题:取给定List中的前TopK个最大的元素并输出。
关键点:
1. topK个最大的元素
2. 我们并不需要顺序,因此一切涉及到sort的工作都是不必要的。
3. 要且只要topK个元素,no more needed!!
因此,我们虽然用堆,但并不需要将整个list都建成堆,只需要维护一个K个元素的堆即可。
其次,我们的堆并不用来排序,这点很重要!!!
具体步骤如下:——求最大的TopK个元素
1. 建一个K个元素的空堆 HP,即一个K个元素的数组
2. 遍历原List,将元素依次压入HP中,在压入过程中:
2.1 如果压入的个数还不到K个,则直接压入,
2.2 如果压入之后的个数达到K个,则将该K个元素维护成一个堆结构
2.3 在压入每个元素之前,如果HP中已经有K个元素,则将新元素与HP的第一个元素比较
2.3.1 如果新元素大于HP的第一个元素,则将新元素放在HP的第一个位置,并调整HP堆。
2.3.2 如果新元素小于HP的第一个元素,则continue 到 2.
遍历完后HP中的K个元素即为List中的前TopK大的元素。
具体代码如下:
heap.h ——heap结构申明
#ifndef HEAP_H #define HEAP_H /* ------------------------------ * Data Define for Heap Structor * Oh This is just a Min Heap!!! * ------------------------------ */ typedef int (*HEAP_CMP_FN)(void *, void*); typedef struct _heap { int len; int size; void ** data; HEAP_CMP_FN heap_cmp_fn; } Heap; /* --------------------------------------------- * Interface Function Define for Heap Operation * --------------------------------------------- */ Heap * heap_create(int size); void * heap_add(Heap * hp, void * pdata); void heap_free(Heap * hp); #endif
heap.c
#include <stdlib.h> #include <stdio.h> #include <string.h> #include "heap.h" /* ---------------------------------- * Default Compare Function for Heap * ---------------------------------- */ int default_heap_cmp_fn(int* m, int* n){ int i = *m - *n; return i > 0 ? 1 : (i < 0 ? -1 : 0); } /* ---------------------------------- * Create and Return a Heap Structor * ---------------------------------- */ Heap * heap_create(int size){ if (size < 2){ fprintf(stderr,"too small size no need heap\n"); return NULL; } Heap * hp = (Heap*)malloc(sizeof(Heap)); hp->len = 0; hp->size = size; hp->data = (void**)malloc(sizeof(void*) * size); memset(hp->data,0,sizeof(void*) * size); hp->heap_cmp_fn = (HEAP_CMP_FN)&default_heap_cmp_fn; return hp; } #define HEAP_ADJUST(hp,l) do { \ int lt = l; \ int i = lt + lt + 1; \ int r = hp->len - 1; \ void * t = hp->data[lt]; \ do { \ if ((i<r) && (hp->heap_cmp_fn(hp->data[i], hp->data[i+1]) > 0)){ \ i += 1; \ } \ if (hp->heap_cmp_fn(t,hp->data[i]) <= 0){ \ break; \ } \ hp->data[lt] = hp->data[i]; \ lt = i; \ i += i + 1; \ }while (i <= r); \ hp->data[lt] = t; \ } while (0); \ /* ---------------------------------- * Add an Element into the Heap * ---------------------------------- */ void * heap_add(Heap * hp, void * pdata){ if (hp->len < hp->size){ hp->data[hp->len] = pdata; hp->len += 1; if (hp->len >= hp->size){ int l = hp->len / 2; while(--l >= 0){ HEAP_ADJUST(hp,l); } } return NULL; } void * top = hp->data[0]; if (hp->heap_cmp_fn(top,pdata) < 0){ hp->data[0] = pdata; HEAP_ADJUST(hp,0); return top; } return pdata; } /* ---------------------------------- * Free the Heap Memory Space * ---------------------------------- */ void heap_free(Heap * hp){ if (hp->data){ for (int i = 0; i < hp->len; i++){ free(hp->data[i]); hp->data[i] = NULL; } free(hp->data); hp->data = NULL; } free(hp); }
heap_test.c
#include <stdlib.h> #include <stdio.h> #include "heap.h" // int main(){ // Heap * hp = heap_create(10); // for (int i = 0; i < 20; i++){ // int c = rand() % 100; // printf("%d,",c); // int * v = (int*)malloc(sizeof(int)); // *v = c; // int * p = heap_add(hp,v); // if (p) // free(p); // } // printf("\n"); // for (int i = 0;i < hp->len;i++){ // printf("%d,",*((int*)hp->data[i])); // } // printf("\n"); // return 0; // } typedef struct { int id; float score; } Element; int comp(Element * m,Element * n){ if (m->score > n->score){ return 1; } else if (m->score<n->score){ return -1; } else{ return 0; } } int main(){ Heap * hp = heap_create(5); hp->heap_cmp_fn = (HEAP_CMP_FN)∁ for (int i = 0; i < 20; i++){ int id = rand() % 100; float score = 100.0 * (float)rand() / RAND_MAX ; Element * p = (Element*)malloc(sizeof(Element)); p->id = id; p->score = score; printf("%d,%f\n",id,score); Element * t = heap_add(hp,p); if (t) free(t); } for (int i = 0;i < hp->len;i++){ fprintf(stderr,"%d,%f\n",((Element *)hp->data[i])->id,((Element *)hp->data[i])->score); } return 0; }
如代码所示,该Heap结构默认支持int类型的数据,如果元素类型为其他,则需要提供一个定制的比较函数,赋值给堆结构的 heap_cmp_fn 函数指针即可,如测试代码中:
hp->heap_cmp_fn = (HEAP_CMP_FN)∁ // comp为定制比较函数