Created
July 1, 2020 16:50
-
-
Save sansyrox/105e37efb01b864bbc99b2338ae91fde to your computer and use it in GitHub Desktop.
Custom Comparators in Python heap
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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