Skip to content
Advertisement

Python priority queue- subsequent pops return wrong values after first pop?

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]

User contributions licensed under: CC BY-SA
3 People found this is helpful
Advertisement