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

点击上方标题跳转至leetcode

题目描述

设计一个找到数据流中第 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 <= 10⁴
0 <= nums.length <= 10⁴
-10⁴ <= nums[i] <= 10⁴
-10⁴ <= val <= 10⁴
最多调用 add 方法 10⁴ 次
题目数据保证，在查找第 k 大元素时，数组中至少有 k 个元素

解法

使用大小为 K 的小根堆，在初始化的时候，保证堆中的元素个数不超过 K 。
在每次 add() 的时候，将新元素 push() 到堆中，如果此时堆中的元素超过了 K，那么需要把堆中的最小元素（堆顶）pop() 出来。
此时堆中的最小元素（堆顶）就是整个数据流中的第 K 大元素。

Python3

class KthLargest:
    def __init__(self, k: int, nums: List[int]):
        self.q = []
        self.size = k
        for num in nums:
            self.add(num)

    def add(self, val: int) -> int:
        heapq.heappush(self.q, val)
        # print(self.q)
        if len(self.q) > self.size:
            heapq.heappop(self.q)
            # print(self.q)
        return self.q[0]
        
nums = [4,5,8,2]
k = 3
kthLargest = KthLargest(k, nums)
param_1 = kthLargest.add(3)
param_2 = kthLargest.add(5)
param_3 = kthLargest.add(10)
param_4 = kthLargest.add(9)
param_5 = kthLargest.add(4)

print(param_1)
print(param_2)
print(param_3)
print(param_4)
print(param_5)

C++

#include<iostream>
#include<vector>
#include<queue>
using namespace std;

class KthLargest {
public:
    priority_queue<int, vector<int>, greater<int>> q;
    int size;
    KthLargest(int k, vector<int>& nums) {
        size = k;
        for (int num : nums) add(num);
    }
    int add(int val) {
        q.push(val);
        if (q.size() > size) q.pop();
        return q.top();
    }
};

int main(){

    vector<int> nums = {4,5,8,2};
    int k = 3;
    KthLargest* obj = new KthLargest(k, nums);
    int param_1 = obj->add(3);
    int param_2 = obj->add(5);
    int param_3 = obj->add(10);
    int param_4 = obj->add(9);
    int param_5 = obj->add(4);
    // int res1 = KthLargest(k,nums).add(3);
    cout << param_1 << "\t" 
         << param_2 << "\t"  
         << param_3 << "\t"  
         << param_4 << "\t"  
         << param_5 << "\t" << endl;
    delete obj;
    return 0;
}
//g++ 703.cpp -std=c++11