Solving the Top k most frequent words problem using a max-heap

Author: aparibocci

Diff for: data_structures/heap/heap.py

@@ -1,9 +1,28 @@
 from __future__ import annotations
 
+from abc import abstractmethod
 from collections.abc import Iterable
+from typing import Generic, Protocol, TypeVar
 
 
-class Heap:
+class Comparable(Protocol):
+    @abstractmethod
+    def __lt__(self: T, other: T) -> bool:
+        pass
+
+    @abstractmethod
+    def __gt__(self: T, other: T) -> bool:
+        pass
+
+    @abstractmethod
+    def __eq__(self: T, other: object) -> bool:
+        pass
+
+
+T = TypeVar("T", bound=Comparable)
+
+
+class Heap(Generic[T]):
     """A Max Heap Implementation
 
     >>> unsorted = [103, 9, 1, 7, 11, 15, 25, 201, 209, 107, 5]
@@ -27,7 +46,7 @@ class Heap:
     """
 
     def __init__(self) -> None:
-        self.h: list[float] = []
+        self.h: list[T] = []
         self.heap_size: int = 0
 
     def __repr__(self) -> str:
@@ -79,7 +98,7 @@ def max_heapify(self, index: int) -> None:
                 # fix the subsequent violation recursively if any
                 self.max_heapify(violation)
 
-    def build_max_heap(self, collection: Iterable[float]) -> None:
+    def build_max_heap(self, collection: Iterable[T]) -> None:
         """build max heap from an unsorted array"""
         self.h = list(collection)
         self.heap_size = len(self.h)
@@ -88,7 +107,7 @@ def build_max_heap(self, collection: Iterable[float]) -> None:
             for i in range(self.heap_size // 2 - 1, -1, -1):
                 self.max_heapify(i)
 
-    def extract_max(self) -> float:
+    def extract_max(self) -> T:
         """get and remove max from heap"""
         if self.heap_size >= 2:
             me = self.h[0]
@@ -102,7 +121,7 @@ def extract_max(self) -> float:
         else:
             raise Exception("Empty heap")
 
-    def insert(self, value: float) -> None:
+    def insert(self, value: T) -> None:
         """insert a new value into the max heap"""
         self.h.append(value)
         idx = (self.heap_size - 1) // 2
@@ -144,7 +163,7 @@ def heap_sort(self) -> None:
     ]:
         print(f"unsorted array: {unsorted}")
 
-        heap = Heap()
+        heap: Heap[int] = Heap()
         heap.build_max_heap(unsorted)
         print(f"after build heap: {heap}")
 

Diff for: strings/top_k_frequent_words.py

@@ -0,0 +1,61 @@
+"""
+Finds the top K most frequent words from the provided word list.
+"""
+
+
+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):
+        self.word = word
+        self.count = count
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, WordCount):
+            return NotImplemented
+        return self.count == other.count
+
+    def __lt__(self, other) -> bool:
+        if not isinstance(other, WordCount):
+            return NotImplemented
+        return self.count < other.count
+
+
+def top_k_frequent_words(words: list[str], k: int) -> list[str]:
+    """
+    Returns the k 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 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']
+    """
+    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, len(count_by_word)))]
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()