原题
实现一个数据结构,提供下面两个接口
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)