algorithms

Algorithm implementations
git clone git://git.laack.co/algorithms.git
Log | Files | Refs | README

top-k-frequent-words-v2.py (825B)

      1 import heapq
      2 
      3 class OrderedWord:
      4     def __init__(self, word):
      5         self.word = word
      6     def __lt__(self, other):
      7         return self.word > other.word
      8 
      9 class Solution:
     10     def topKFrequent(self, words: List[str], k: int) -> List[str]:
     11         
     12         counts = {}
     13         for word in words:
     14             if word in counts:
     15                 counts[word] += 1
     16             else:
     17                 counts[word] = 1
     18         
     19         heap = []
     20 
     21         for word in counts:
     22             heappush (heap, (counts[word], OrderedWord(word)))
     23             if len(heap) > k:
     24                 heappop (heap)
     25 
     26         heap_ordered = []
     27 
     28         for item in heap:
     29             heappush (heap_ordered, (-item[0], item[1].word))
     30     
     31         res = []
     32 
     33         for i in range(k):
     34             res.append(heappop(heap_ordered)[1])
     35         return res