Leetcode——703. 数据流中的第K大元素

题目描述

设计一个找到数据流中第K大元素的类（class）。注意是排序后的第K大元素，不是第K个不同的元素。

你的 KthLargest 类需要一个同时接收整数 k 和整数数组nums 的构造器，它包含数据流中的初始元素。每次调用 KthLargest.add，返回当前数据流中第K大的元素。

示例:

int k = 3;
int[] arr = [4,5,8,2];
KthLargest kthLargest = new KthLargest(3, arr);
kthLargest.add(3);   // returns 4
kthLargest.add(5);   // returns 5
kthLargest.add(10);  // returns 5
kthLargest.add(9);   // returns 8
kthLargest.add(4);   // returns 8

说明:
你可以假设 nums 的长度≥ k-1 且k ≥ 1。

解题思路

• 使用python的标准库的heapq模块构建一个最小堆，最小堆的大小为k
• 最小堆里面放的元素就是前k大的元素
• 每进来一个元素，就和这个对的最小值比较，如果大于这个最小值，就替换

代码

class KthLargest(object):

    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        import heapq
        self.heapq = nums 
        self.size = len(self.heapq) # 当前堆的大小
        self.k = k # 前k大的元素的长度 
        heapq.heapify(self.heapq) # 将列表转换为堆
        while self.size > k:
            heapq.heappop(self.heapq) # 弹出小于第k大的元素，保留前k大的元素
            self.size -= 1  # 堆的大小减1

        
    def add(self, val):
        """
        :type val: int
        :rtype: int
        """
        if self.size < self.k: # 如果当前堆的大小小于k
            heapq.heappush(self.heapq,val)  # 将新值添加到最小堆
            self.size += 1 # 堆的大小加一
        elif val > self.heapq[0]: # 如果最小堆的最小元素比新值 小
            heapq.heapreplace(self.heapq,val) # 用新值替换
        return self.heapq[0]  # 返回堆的根节点元素即为第k大的元素
        

参考资料

python堆排序heapq：https://www.jianshu.com/p/e003872fa7b9