-
Notifications
You must be signed in to change notification settings - Fork 1
/
heap.py
29 lines (26 loc) · 839 Bytes
/
heap.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import collections
import heapq
class Solution:
# Time Complexity = O(n + nlogk)
# Space Complexity = O(n)
def topKFrequent(self, words, k):
count = collections.Counter(words)
heap = []
for key, value in count.items():
heapq.heappush(heap, Word(value, key))
if len(heap) > k:
heapq.heappop(heap)
res = []
for _ in range(k):
res.append(heapq.heappop(heap).word)
return res[::-1]
class Word:
def __init__(self, freq, word):
self.freq = freq
self.word = word
def __lt__(self, other):
if self.freq == other.freq:
return self.word > other.word
return self.freq < other.freq
def __eq__(self, other):
return self.freq == other.freq and self.word == other.word