top_k.py

import heapq
from collections import deque

    
class TopK():

    def __init__(self, k=5):
        self.k = k
        self._initialize()

    def _initialize(self):
        self.ranking = []
        heapq.heapify(self.ranking)
        self.dq = deque(maxlen=self.k)
        self.num = 0

    def add(self, key, val):
        if key not in self.dq:
            if self.num < self.k:
                heapq.heappush(self.ranking, (val, key))
                self.dq.append(key)
                self.num += 1
            else:
                (min_val, min_key) = heapq.heappushpop(self.ranking, (val, key))
                if key != min_key:
                    self.dq.remove(min_key)
                    self.dq.append(key)

    def most_common(self, k=0):
        k = k or self.k
        for (val, key) in heapq.nlargest(k, self.ranking):
            yield (key, val)