src/sts/projects/SWINTS/swints/topk.py

import heapq

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

    def reset(self):
        self.data = []

    def size(self):
        return len(self.data)

    def push(self, x):
        if len(self.data) < self.k:
            heapq.heappush(self.data, x)
        else:
            heapq.heappushpop(self.data, x)

    def extract(self, sort=False):
        if sort:
            self.data.sort(reverse=True)
        return self.data