创建小根堆, 使用python内部的模块 heapq
import heapq
import randomclass TopK:def __init__(self, li, k):self.minheap = []self.capacity = kself.li = lidef push(self, val):if len(self.minheap) >= self.capacity:# 大于等于想要的元素个数min_val = self.minheap[0]if val < min_val: # 新元素比最小的还小, 直接忽略passelse:# 比最小的大, 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.minheapdef 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]