finds the top k most frequent words from the provided word list this implementation aims to show how to solve the problem using the heap class already present in this repository computing order statistics is in fact a typical usage of heaps this is mostly shown for educational purposes since the problem can be solved in a few lines using collections counter from the python standard library from collections import counter def topkfrequentwordswords kvalue return x0 for x in counterwords mostcommonkvalue wordcount a 1 eqwordcount b 1 true wordcount a 1 eqwordcount a 1 true wordcount a 1 eqwordcount a 2 false wordcount a 1 eqwordcount b 2 false wordcount a 1 eq1 notimplemented wordcount a 1 ltwordcount b 1 false wordcount a 1 ltwordcount a 1 false wordcount a 1 ltwordcount a 2 true wordcount a 1 ltwordcount b 2 true wordcount a 2 ltwordcount a 1 false wordcount a 2 ltwordcount b 1 false wordcount a 1 lt1 notimplemented returns the kvalue most frequently occurring words in nonincreasing order of occurrence in this context a word is defined as an element in the provided list in case kvalue is greater than the number of distinct words a value of k equal to the number of distinct words will be considered instead topkfrequentwords a b c a c c 3 c a b topkfrequentwords a b c a c c 2 c a topkfrequentwords a b c a c c 1 c topkfrequentwords a b c a c c 0 topkfrequentwords 1 topkfrequentwords a a 2 a wordcount a 1 __eq__ wordcount b 1 true wordcount a 1 __eq__ wordcount a 1 true wordcount a 1 __eq__ wordcount a 2 false wordcount a 1 __eq__ wordcount b 2 false wordcount a 1 __eq__ 1 notimplemented wordcount a 1 __lt__ wordcount b 1 false wordcount a 1 __lt__ wordcount a 1 false wordcount a 1 __lt__ wordcount a 2 true wordcount a 1 __lt__ wordcount b 2 true wordcount a 2 __lt__ wordcount a 1 false wordcount a 2 __lt__ wordcount b 1 false wordcount a 1 __lt__ 1 notimplemented returns the k_value most frequently occurring words in non increasing order of occurrence in this context a word is defined as an element in the provided list in case k_value is greater than the number of distinct words a value of k equal to the number of distinct words will be considered instead top_k_frequent_words a b c a c c 3 c a b top_k_frequent_words a b c a c c 2 c a top_k_frequent_words a b c a c c 1 c top_k_frequent_words a b c a c c 0 top_k_frequent_words 1 top_k_frequent_words a a 2 a
from collections import Counter
from functools import total_ordering

from data_structures.heap.heap import Heap


@total_ordering
class WordCount:
    def __init__(self, word: str, count: int) -> None:
        self.word = word
        self.count = count

    def __eq__(self, other: object) -> bool:
        
        if not isinstance(other, WordCount):
            return NotImplemented
        return self.count == other.count

    def __lt__(self, other: object) -> bool:
        
        if not isinstance(other, WordCount):
            return NotImplemented
        return self.count < other.count


def top_k_frequent_words(words: list[str], k_value: int) -> list[str]:
    
    heap: Heap[WordCount] = Heap()
    count_by_word = Counter(words)
    heap.build_max_heap(
        [WordCount(word, count) for word, count in count_by_word.items()]
    )
    return [heap.extract_max().word for _ in range(min(k_value, len(count_by_word)))]


if __name__ == "__main__":
    import doctest

    doctest.testmod()