自定义比较函数的Heap类

参考:https://www.cnblogs.com/yaozhaoyz/p/3595842.html

import heapq
import random

class Element():
    def __init__(self, a, b, c):
        self.a = a
        self.b = b
        self.c = c

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
        
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))

猜你喜欢

转载自blog.csdn.net/weixin_40876685/article/details/88972252
今日推荐