创建小根堆, 使用python内部的模块 heapq
import heapq
import random
class TopK:
def __init__(self, li, k):
self.minheap = []
self.capacity = k
self.li = li
def push(self, val):
if len(self.minheap) >= self.capacity:
# 大于等于想要的元素个数
min_val = self.minheap[0]
if val < min_val:
# 新元素比最小的还小, 直接忽略
pass
else:
# 比最小的大, name进行替换, 先取出最小的, 然后添加新元素, 内部会重新调整小堆根
heapq.heapreplace(self.minheap, val)
else:
# 个数小于k, 直接添加
heapq.heappush(self.minheap, val)
def get_topk(self):
for val in self.li:
self.push(val)
return self.minheap
def test():
li = list(range(1000))
random.shuffle(li)
obj = TopK(li, 10)
print(obj.get_topk())
if __name__ == '__main__':
test()
结果:
[990, 992, 991, 993, 995, 994, 996, 997, 998, 999]