Python中自带的堆heapq,不支持自定义的比较函数。 这导致,heapq中的元素,如果是结构体的话,不太方便。
实现了一个支持自定义比较函数的Heap类。
import heapq
import random
class MyHeap(object):
def __init__(self, initial=None, key=lambda x:x):
self.k = 20 # the Size of this Heap
self.key = key
self._data = []
def push(self, item):
if len(self._data) < self.k:
heapq.heappush(self._data, (self.key(item), item))
else:
topk_small = list(self._data[0])
if item.a > topk_small[1].a:
heapq.heapreplace(self._data, (self.key(item), item))
def pop(self):
if(len(self._data)>1):
return heapq.heappop(self._data)[1]
else:
return None
class Element():
def __init__(self, a,b,c):
self.a = a
self.b = b
self.c = c
myHeap = MyHeap(key=lambda item:item.a)
for i in range(100):
a = random.randint(0,100);
b = random.randint(0,100);
myHeap.push(Element(a,b,b))
for i in range(20):
j = myHeap.pop()
if(j!=None):
print ""+str(j.a) + "\t" + str(j.b)