python使用heapq实现小顶堆（TopK大）/大顶堆（BtmK小）

朴灿烈づ我的快乐病毒、 2022-07-12 07:44 282阅读 0赞

参考链接

[https://www.coder4.com/archives/3844][https_www.coder4.com_archives_3844]

求一个数列前K大数的问题经常会遇到，在程序中一般用小顶堆可以解决，下面的代码是使用python的heapq实现的小顶堆示例代码：

# !/usr/bin/env python
        # -*- coding:gbk -*-
    
        import sys
        import heapq
    
        class TopKHeap(object):
            def __init__(self, k):
                self.k = k
                self.data = []
    
            def push(self, elem):
                if len(self.data) < self.k:
                    heapq.heappush(self.data, elem)
                else:
                    topk_small = self.data[0]
                    if elem > topk_small:
                        heapq.heapreplace(self.data, elem)
            def topk(self):
                return [x for x in reversed([heapq.heappop(self.data) for x in xrange(len(self.data))])]
    
    
        def main():
            list_num = [1, 2, 3, 4, 5, 6, 7, 8, 9]
            th = TopKHeap(5)
    
            for i in list_num:
                th.push(i)
    
            print th.topk()
    
        if __name__ == "__main__":
            main()

python的heapq在实现的时候，没有像STL或者Java可以传入比较函数，具体的原因可以参考参考文档给出的链接。

因此有些人想出了比较trick的思路。一句话概括如下：

push(e)改为push(-e)，pop(e)为-pop(e)，也就是说存入和取出的数都是相反数，其他逻辑和TopK相同。（点赞）

实现用户自定义的比较函数，允许elem是一个tuple，按照tuple的第一个元素进行比较，所以可以把tuple的第一个元素作为我们的比较的key。

英文原话：

> The heapq documentation suggests that heap elements could be tuples in which the first element is the priority and defines the sort order.

import heapq
        class MyHeap(object):
           def __init__(self, initial=None, key=lambda x:x):
               self.key = key
               if initial:
                   self._data = [(key(item), item) for item in initial]
                   heapq.heapify(self._data)
               else:
                   self._data = []
    
           def push(self, item):
               heapq.heappush(self._data, (self.key(item), item))
    
           def pop(self):
               return heapq.heappop(self._data)[1]

[https_www.coder4.com_archives_3844]: https://www.coder4.com/archives/3844