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:
from queue import PriorityQueue class Solution: def topKFrequent(self, nums: List[int], k: int) -> List[int]: counts = {} mynums = set() q = PriorityQueue() for num in nums: if num not in counts: print(str(num) + " not in counts, setting count to 1") counts[num] = 1 mynums.add(num) else: print(str(num) + " in counts, adding 1 to count") counts[num] = counts[num] + 1 print(str(num) + " count is now = " +str(counts[num])) for num in mynums: print("enqueing " + str(num) + " with count " + str(counts[num])) q.put(num, -1*counts[num]) ans = [] while k > 0: cur = q.get() print("adding " + str(cur) + " to ans") ans.append(cur) k = k - 1 return ans
Here is stdout when running this with k = 2 and nums = [4,1,-1,2,-1,2,3]
4 not in counts, setting count to 1 1 not in counts, setting count to 1 -1 not in counts, setting count to 1 2 not in counts, setting count to 1 -1 in counts, adding 1 to count -1 count is now = 2 2 in counts, adding 1 to count 2 count is now = 2 3 not in counts, setting count to 1 enqueing 1 with count 1 enqueing 2 with count 2 enqueing 3 with count 1 enqueing 4 with count 1 enqueing -1 with count 2 adding -1 to ans adding 1 to ans
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]