本文主要是介绍python 小顶堆，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！
转自http://blog.csdn.net/tanghaiyu777/article/details/55271004
参考链接
https://www.coder4.com/archives/3844
求一个数列前K大数的问题经常会遇到，在程序中一般用小顶堆可以解决，下面的代码是使用python的heapq实现的小顶堆示例代码：
import sysimport heapqclass TopKHeap(object):def __init__(self, k):self.k = kself.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()
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
- 25
- 26
- 27
- 28
- 29
- 30
- 31
- 32
- 33
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 heapqclass MyHeap(object):def __init__(self, initial=None, key=lambda x:x):self.key = keyif 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]
这篇关于python 小顶堆的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！
