I am trying to solve Leetcode #347 Top K Frequent elements using a priority queue. However, my code is failing some test cases.
Here is the code:
JavaScript
x
27
27
1
from queue import PriorityQueue
2
class Solution:
3
def topKFrequent(self, nums: List[int], k: int) -> List[int]:
4
counts = {}
5
mynums = set()
6
q = PriorityQueue()
7
for num in nums:
8
if num not in counts:
9
print(str(num) + " not in counts, setting count to 1")
10
counts[num] = 1
11
mynums.add(num)
12
else:
13
print(str(num) + " in counts, adding 1 to count")
14
counts[num] = counts[num] + 1
15
print(str(num) + " count is now = " +str(counts[num]))
16
for num in mynums:
17
print("enqueing " + str(num) + " with count " + str(counts[num]))
18
q.put(num, -1*counts[num])
19
ans = []
20
while k > 0:
21
cur = q.get()
22
print("adding " + str(cur) + " to ans")
23
ans.append(cur)
24
k = k - 1
25
return ans
26
27
Here is stdout when running this with k = 2 and nums = [4,1,-1,2,-1,2,3]
JavaScript
1
17
17
1
4 not in counts, setting count to 1
2
1 not in counts, setting count to 1
3
-1 not in counts, setting count to 1
4
2 not in counts, setting count to 1
5
-1 in counts, adding 1 to count
6
-1 count is now = 2
7
2 in counts, adding 1 to count
8
2 count is now = 2
9
3 not in counts, setting count to 1
10
enqueing 1 with count 1
11
enqueing 2 with count 2
12
enqueing 3 with count 1
13
enqueing 4 with count 1
14
enqueing -1 with count 2
15
adding -1 to ans
16
adding 1 to ans
17
Why does my code seem to skip over 2 as the the entry to return on the second call to q.get()?
It correctly returns -1 on the first time, but the second pop it returns 1 and not 2. I cannot find the reasoning for this. Please help!
Advertisement
Answer
the PriorityQue object does not sort items as your code implies you think. What you want to do instead is:
1- change q.put(num, -1*counts[num])
to q.put((-1 * counts[num], num))
2- change cur = q.get()
to cur = q.get()[1]