update Python/Top_K_Frequent_Elements.py

Author: kumailn

Python/Top_K_Frequent_Elements.py

@@ -1,9 +1,10 @@
 # Question(#347): Given a non-empty array of integers, return the k most frequent elements.
 # Solution: Use a counting sort to place the most frequent items into buckets, loop through the buckets backwards,
 #           (since the last bucket cooresponds to the most common item), add the items in the bucket to our result and decrement
-#           value of k by the amount of items in the bucket
+#           value of k by the amount of items in the bucket. An alternate solution is to use a Heap.
 
-def topKFrequent(self, nums: List[int], k: int) -> List[int]:
+# Using a counting sort
+def topKFrequent(nums: List[int], k: int) -> List[int]:
     # Make an array with len(nums) + 1 positions, because the max frequency of any element 
     # can coorespond to the len(nums) + 1 index.
     buckets = [[] for _ in range(len(nums) + 1)]
@@ -23,4 +24,21 @@ def topKFrequent(self, nums: List[int], k: int) -> List[int]:
             result += i
             k -= len(i)
 
+    return result
+
+# Using a max heap
+def topKFrequent(nums: List[int], k: int) -> List[int]:
+    result, heap = [], []
+
+    # Add every item to the max-heap, (Note we add -b instead of b since python supports min-heaps),
+    # each item consists of its count and its value. 
+    for a, b in collections.Counter(nums).items(): 
+        heapq.heappush(heap, (-b, a))
+
+    # As long as we either have items in the heap and we have k left, we want to take things
+    # off of the max heap. Pop the top, add the value associated with it to the result and decrement k.
+    while k and heap:
+        result += [heapq.heappop(heap)[1]]
+        k -= 1
+        
     return result