14_堆TopK排序

import time
import heapq
import random

class TopK:
    """
    获取k个元素，固定内存
    思路：
    1 放入元素前，建立一个最小堆，k个元素
    2 迭代剩余元素-若当前元素小于堆顶元素，跳过；否则替换堆顶元素为当前元素，并重新调整堆
    """
    def __init__(self, iterable, k):
        self.minheadp = []
        self.capacity = k
        self.iterable = iterable

    def push(self, val):
        if len(self.minheadp) >= self.capacity:
            min_val = self.minheadp[0]
            if val < min_val:
                pass
            else:
                # 返回并删除堆中的最小item
                # 返回并且pop堆顶最小值，推入新的 val 值并调整堆
                heapq.heapreplace(self.minheadp, val)
        else:
            # 前面 k 个元素直接放入minheap
            heapq.heappush(self.minheadp, val)

    def get_topk(self):
        for val in self.iterable:
            self.push(val)
        return self.minheadp

def test():
    i = list(range(100)) # 可迭代元素，节省内存
    random.shuffle(i)
    _ = TopK(i, 10)
    print(_.get_topk())

if __name__ == '__main__':
    test()