python中的堆支持自定义的比较函数 - Heap in Python with comparator.

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)