LintCode 545 [Top k Largest Number II]

原题

实现一个数据结构，提供下面两个接口
1.add(number) 添加一个元素
2.topk() 返回前K大的数

样例

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]

解题思路

• 内部维护一个大小为k的最小堆
• 如果size < k则直接将num加入到minHeap
• 如果size >= k, 则比较num与minHeap中的最小值，若小于最小值则忽略，大于最小值则删除最小值，并把num加入到minHeap
• 每次返回minHeap数组，返回时先反向排序

sorted(nums, reverse=True)

完整代码

import Queue

class Solution:

    # @param {int} k an integer
    def __init__(self, k):
        # initialize your data structure here.
        self.size = k
        self.minHeap = Queue.PriorityQueue()

        
    # @param {int} num an integer
    def add(self, num):
        # Write your code here
        if self.minHeap.qsize() < self.size:
            self.minHeap.put(num)
        elif num > self.minHeap.queue[0]:
            self.minHeap.get()
            self.minHeap.put(num)


    # @return {int[]} the top k largest numbers
    def topk(self):
        # Write your code here
        return sorted(self.minHeap.queue, reverse=True)