Skip to content

Instantly share code, notes, and snippets.

@sansyrox
Created July 1, 2020 16:50
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save sansyrox/105e37efb01b864bbc99b2338ae91fde to your computer and use it in GitHub Desktop.
Save sansyrox/105e37efb01b864bbc99b2338ae91fde to your computer and use it in GitHub Desktop.
Custom Comparators in Python heap
class FreqWord(object):
def __init__(self, freq, word):
self.freq = freq
self.word = word
def __lt__(self, other):
if self.freq != other.freq:
return lt(self.freq, other.freq)
else:
# opposite sort
return lt(other.word, self.word)
class Solution:
def topKFrequent(self, words: List[str], k: int) -> List[str]:
words_with_count = collections.Counter(words)
min_heap = list()
for word, count in words_with_count.items():
heapq.heappush(min_heap, FreqWord(count, word))
if len(min_heap) > k:
heapq.heappop(min_heap)
return [heapq.heappop(min_heap).word for _ in range(k)][::-1]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment