python利用heapq实现小顶堆(查找最大的N个元素)

查找最大的N个元素——堆数据结构

给出序列,求出TopK大的元素,使用小顶堆,heapq模块实现

使用heappush()时,从数据源增加新元素时会保持元素的堆顺序。
在一个操作中删除现有元素并替换为新值,可以使用heapreplace()

import heapq

class BtmkHeap(object):
def __init__(self, k):
self.k = k
self.data = []

def Push(self, elem):
# Reverse elem to convert to max-heap
num1 = elem[0]
num1 = -num1
num2 = elem[1]
# Using heap algorighem
if len(self.data) < self.k:
heapq.heappush(self.data, (num1,num2))
else:
topk_small = self.data[0][0]
if elem[0] > topk_small:
heapq.heapreplace(self.data,(num1,num2))

def BtmK(self):
return sorted([(-x[0],x[1]) for x in self.data])


if __name__ == '__main__':
a = (100,1)
b = (200,2)
c = (300,3)
d = (400,4)
e = (50,5)
btm = BtmkHeap(3)
btm.Push(a)
btm.Push(b)
btm.Push(c)
btm.Push(e)
btm.Push((23,1))
print(btm.BtmK())



结果[(23, 1), (50, 5), (100, 1)]

posted @ 2018-03-10 15:56  circle元  阅读(1837)  评论(0编辑  收藏  举报