Lintcode 545. Top k Largest Numbers II (Medium) (Python)

Top k Largest Numbers II

Description:

Implement a data structure, provide two interfaces:

add(number). Add a new number in the data structure.
topk(). Return the top k largest numbers in this data structure. k is given when we create the data structure.
Example
s = new Solution(3);

create a new data structure.
s.add(3)
s.add(10)
s.topk()
return [10, 3]
s.add(1000)
s.add(-99)
s.topk()
return [1000, 10, 3]
s.add(4)
s.topk()
return [1000, 10, 4]
s.add(100)
s.topk()
return [1000, 100, 10]

Code:

import Queue
class Solution:
    """
    @param: k: An integer
    """
    def __init__(self, k):
        # do intialization if necessary
        self.size = k
        self.heap = Queue.PriorityQueue()
    """
    @param: num: Number to be added
    @return: nothing
    """
    def add(self, num):
        # write your code 
        if self.heap.qsize() < self.size:
            self.heap.put(num)
        elif num > self.heap.queue[0]:
            self.heap.get()
            self.heap.put(num)

    """
    @return: Top k element
    """
    def topk(self):
        # write your code here
        return sorted(self.heap.queue, reverse=True)