利用最小堆实现topk

import heapq

class TopK:
    '''
    思路：
    1.先放入元素前k个建立一个最小堆
    2.迭代剩余元素：
    如果当前元素小于堆顶元素，跳过该元素
    否则替换堆顶为当前元素，并重新调整堆

    用最小堆来获取最大的几个元素,用前面最小的值替换成最大的值
    '''
    def __init__(self,iterable,k):
        self.minheap = [] #最小堆
        self.capacity = k # 最大的几个数
        self.iterable = iterable # 传入的列表
    def push(self,val): # 传入数据
        # 如果超过k个元素就push
        if len(self.minheap) >= self.capacity:
            min_val = self.minheap[0]
            if val < min_val:
                pass # 忽略元素
            else:
                # 替换最小堆元素，并保持最小堆的性质
                heapq.heapreplace(self.minheap,val)
        else:
            # 没有超过k个元素则push
            heapq.heappush(self.minheap,val)
    def get_topk(self):
        for val in self.iterable:
            self.push(val) # 出入数据
        return self.minheap # 返回最小堆

def test():
    import random
    i = list(range(1000))
    random.shuffle(i)
    _ = TopK(i,3)
    print(_.get_topk())
test()