leet703.kth-largest-element-in-a-stream

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

请实现 KthLargest 类：

• KthLargest(int k, int[] nums) 使用整数 k 和整数流 nums 初始化对象。
• int add(int val) 将 val 插入数据流 nums 后，返回当前数据流中第 k 大的元素。

示例 1：

输入：
[“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); // 返回 4
kthLargest.add(5); // 返回 5
kthLargest.add(10); // 返回 5
kthLargest.add(9); // 返回 8
kthLargest.add(4); // 返回 8

示例 2：

输入：
[“KthLargest”, “add”, “add”, “add”, “add”]
[[4, [7, 7, 7, 7, 8, 3]], [2], [10], [9], [9]]

输出：[null, 7, 7, 7, 8]

解释：

KthLargest kthLargest = new KthLargest(4, [7, 7, 7, 7, 8, 3]);
kthLargest.add(2); // 返回 7
kthLargest.add(10); // 返回 7
kthLargest.add(9); // 返回 7
kthLargest.add(9); // 返回 8

提示：

• 0 <= nums.length <= 10^4
• 1 <= k <= nums.length + 1
• -10^4 <= nums[i] <= 10^4
• -10^4 <= val <= 10^4
• 最多调用 add 方法 10^4 次

---

思路

这是堆的经典应用。第 k 大/小问题。

• 一堆数，求第 k 大：建立大小为 k 的小顶堆。
• 求第 k 小：建立大小为 k 的大顶堆。

本题求第 k 大，则建立大小为 k 的小顶堆。

Code

class KthLargest {
public:
    priority_queue<int, vector<int>, greater<>> pq;
    int k;
    
    KthLargest(int k, vector<int>& nums) {
        this->k = k;
        for (int e : nums)
            add(e);
    }
    
    int add(int val) {
        pq.push(val);
        if (pq.size() > k)
            pq.pop();
        return pq.top();
    }
};
 
/**
 * Your KthLargest object will be instantiated and called as such:
 * KthLargest* obj = new KthLargest(k, nums);
 * int param_1 = obj->add(val);
 */

拓展

上面的写法很clean，因为只用存 k 个元素。在 add 的过程中，没有入堆的全部不用考虑，因为将来也不可能入堆。如果再让你设计一个 pop 函数，支持弹出元素呢？使用下面的写法，可以处理 pop 的情况。

class KthLargest {
public:
    int k;
    vector<int> nums;
    KthLargest(int k, vector<int>& nums) {
        this->k = k;
        this->nums = nums;
        const int n = nums.size();
        if (n >= 2) {
            // 先将数组从大到小排序，然后用前面k个大元素建立小顶堆
            // 则堆顶就是第k大元素。需要考虑 k = n + 1 的边界情况
            sort(this->nums.begin(), this->nums.end(), greater<>());
            make_heap(this->nums.begin(), this->nums.begin() + min(n, k), greater<>());
        }
    }
    
    int add(int val) {
        nums.push_back(val);
        if (k >= nums.size()) { // 根据题目约束，k至多等于push一个元素后的nums.size()
            push_heap(nums.begin(), nums.end(), greater<>());
            return nums.front();
        }
        if (val > nums.front()) { // val需要入堆
            // 把val换到堆尾后面一个位置
            swap(nums[k], nums.back());
            // 弹出原堆顶的同时，将val也入堆了。这里利用了pop_heap的操作，注意原来heap size是k。
            // 这里多算一个nums[k]，就是为了在弹出堆顶的同时，把nums[k]=val换到堆顶进行下沉。
            pop_heap(nums.begin(), nums.begin() + k + 1, greater<>());
 
        }
        return nums.front();
    }
};