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

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

题干

设计一个找到数据流中第 k 大元素的类（class）。注意是排序后的第 k 大元素，不是第 k 个不同的元素。
请实现 KthLargest 类：
KthLargest(int k, int[] nums) 使用整数 k 和整数流 nums 初始化对象。
int add(int val) 将 val 插入数据流 nums 后，返回当前数据流中第 k 大的元素。

示例：
输入：
[“KthLargest”, “add”, “add”, “add”, “add”, “add”]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
输出：
[null, 4, 5, 5, 8, 8]

解释：
KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);
kthLargest.add(3); // return 4
kthLargest.add(5); // return 5
kthLargest.add(10); // return 5
kthLargest.add(9); // return 8
kthLargest.add(4); // return 8

提示：
1 <= k <= 104
0 <= nums.length <= 104
-104 <= nums[i] <= 104
-104 <= val <= 104
最多调用 add 方法 104 次
题目数据保证，在查找第 k 大元素时，数组中至少有 k 个元素

题解

仔细一想，题目求得是第k“大”的元素，那其实只要第一次找到第k大的元素后，比这个数小的数都可以考虑不计算，因为根本不影响答案。由此，也就不需要对整个数组进行排序了，可以换一种更高效的数据结构。
这里用了优先队列实现小顶堆，初始化时先将比第k大的数小的数全部抛掉，使得堆顶即为第k大元素。
在加入新的数时，只需与堆顶元素进行比较，如果小于堆顶元素则不影响答案，直接返回堆顶元素即可。如果大于堆顶元素，则需要更新答案，将元素加入优先队列，然后抛掉原先的堆顶元素。
特别需要注意的是，有可能在第一次调用add方法时优先队列中的元素还不足k个，此时不论元素大小直接加入即可，不需要抛出操作。

class KthLargest {
public:
    int kth;
    priority_queue<int,vector<int>,greater<int> > kthNum;
    KthLargest(int k, vector<int>& nums) {
        kth = k;
        for(auto i : nums){
            kthNum.push(i);
        }
        int popCount = nums.size() - k;
        for(int i = 0 ; i < popCount ; ++i){
            kthNum.pop();
        }
    }
    
    int add(int val) {
        if(kthNum.size() < kth){
            kthNum.push(val);
            return kthNum.top();
        }
        if(val > kthNum.top()){
            kthNum.push(val);
            kthNum.pop();
        }
        return kthNum.top();
    }
};