Top K elements

692. Top K Frequent Words

N log k solution is here

from functools import total_ordering
@total_ordering
class Element:
    def __init__(self, word, freq):
        self.word, self.freq = word, freq
    def __eq__(self, other):
        return self.freq==other.freq and self.word==other.word
    def __lt__(self, other):
        if self.freq==other.freq:
            return self.word>other.word
        return self.freq<other.freq
    
class Solution:
    def topKFrequent(self, words: List[str], k: int) -> List[str]:
        cnt = collections.Counter(words)
        hq = []
        for word, freq in cnt.items():
            if len(hq)==k:
                f =  heapq.heappushpop
            else:
                f = heapq.heappush 
            f(hq,Element(word, freq))
        res = []
        while hq:
            res.append(heapq.heappop(hq).word)
        return res[::-1]

Reference

[1] https://leetcode.com/problems/top-k-frequent-words/discuss/108348/Python-3-solution-with-O(nlogk)-and-O(n)