python heapq topk

foo.py

import heapq


class TopK(object):
    def __init__(self, k):
        self.k = k
        self.heap = []

    def push(self, elem):
        if len(self.heap) < self.k:
            heapq.heappush(self.heap, elem)
        elif elem > self.heap[0]:
            heapq.heapreplace(self.heap, elem)

    def topk(self):
        return list(heapq.heappop(self.heap) for _ in range(len(self.heap)))[::-1]


class BtmK(object):
    def __init__(self, k):
        self.k = k
        self.heap = []

    def push(self, elem):
        if len(self.heap) < self.k:
            heapq.heappush(self.heap, -elem)
        elif elem < -self.heap[0]:
            heapq.heapreplace(self.heap, -elem)

    def btmk(self):
        return list(-heapq.heappop(self.heap) for _ in range(len(self.heap)))[::-1]


if __name__ == '__main__':
    nums = [1, 7, 4, 43, 11, 5, 54, 42, 3, 13, 33]

    top4 = TopK(4)
    for num in nums:
        top4.push(num)
    print(top4.topk())

    btm4 = BtmK(4)
    for num in nums:
        btm4.push(num)
    print(btm4.btmk())

"""
ref

https://docs.python.org/3.5/library/heapq.html
https://www.coder4.com/archives/3844
https://blog.csdn.net/tanghaiyu777/article/details/55271004
https://www.jianshu.com/p/801318c77ab5

"""