需求是这样的,我们都知道,在信息检索中,经常要取top-k(一共k,而不是第k)个得分最大的文档,并且从大到小排序。
而且文档规模很大,最少也要上千万。
话说这是一道很可以拿来面试的题啊。
我们不考虑Hadoop神马的,就说说单机怎么搞。
最傻的做法就是把1000万个都存储下来,然后sort,然后取min(k, vec.size())。
这样有两个缺点:
1、内存占用非常大,其实我们只要保留最大的1000个,但这样就要保存N个。在1000万的测试中,它要占用68M内存之多。
2、速度慢虽然是快速排序,但我们实际只要前1000个,后面那些排序都是无用功了。
于是靠谱的做法就是构建一个小顶堆:
(1)当堆中元素数量小于k时, 向堆中加入元素并调整堆
(2)当堆中元素数量大于等于k时 ,将堆顶(显然是当前最小的)和新元素比较:
(a)如果新元素小,丢弃之,我们只要top k 最大的。
(b) 如果新元素比堆顶大,把堆顶用新元素换掉,然后调整堆。
好了,理论是上面这个过程,我们一般就不要裸写堆了,用c++ stl的pop_heap/make_heap/push_heap搞定之。注意它没有调整堆单独的函数,而是第一次make_heap,之后用成对的pop_heap/push_heap来完成堆维护(这两个操作不会改变堆大小)。而堆是构建在vector基础上的。
细节是:上述默认都是构造大顶堆,从小到大排序等。于是我们要子构造一个反向排序函数。
好了,再重复一下堆:
小顶堆:堆顶始终是最小的元素,排序时把最小的换到末尾,类似选择排序的过程。这样拍出来的是从大到小。
大顶堆:排序之后从小到大。
比较一下性能(1kw个数据,值域0~1kw):
算法/内存/时间:
(1)sort topk / 68MB / 11s
(2)heap / 3MB / 6s
内存优势非常明显。时间优势在这个规模可能还显示不出来,也可能是qsort的随机性能。
堆解决代码:
#include <iostream> #include <algorithm> #include <vector> using namespace std; //type use long typedef long TYPE; //反向比较器,将make_heap等函数转化为小顶堆 bool rcmp (const TYPE &a,const TYPE &b) { return a>b; } //An vector for global, would keep at most k elements vector<TYPE> vec; TYPE k; //recv elemnts and keep top K(big), select top k void recv_keep_top_k(TYPE elem) { if(vec.size()<k) { vec.push_back(elem); if(vec.size()>=k) { make_heap(vec.begin(), vec.end(), rcmp); } } else { // exceed k, need select //compare to smallest if(elem < vec.front() ) { //small than heap's small, drop return; } //bigger than smallest in heap, keep and replace the small one pop_heap(vec.begin(), vec.end(), rcmp); vec.back() = elem; push_heap(vec.begin(), vec.end(), rcmp); return; } } int main () { k = 100; TYPE e; while(cin>>e) { recv_keep_top_k(e); } //sort from small to big //Each time, the smallest one was swap to tail if(vec.size()< k) { make_heap(vec.begin(), vec.end(), rcmp); } sort_heap(vec.begin(),vec.end(), rcmp); cout << "Final top " << k << endl; for (unsigned i=0; i<vec.size(); i++) { cout << vec[i] << endl; } return 0; }
对照组小白鼠(sort topK):
#include <iostream> #include <algorithm> #include <vector> using namespace std; //type use long typedef long TYPE; bool rcmp (const TYPE &a,const TYPE &b) { return a>b; } //An vector for global, would keep at most k elements vector<TYPE> vec; TYPE k; int main () { k = 100; TYPE e; while(cin>>e) { vec.push_back(e); } //sort from big to small sort(vec.begin(), vec.end(), rcmp); cout << "Final top " << k << endl; for (unsigned i=0; i<k && i<vec.size(); i++) { cout << vec[i] << endl; } return 0; }
构建在 vector 上的堆会不会比较慢?比如从堆中拿掉一个元素这种操作会不会导致 vector 里面的后续元素都被拷贝一遍?如果自己用链表+二分查找模拟一个堆应该是最快了吧。