题目链接
设计一个找到数据流中第 k 大元素的类（class）。注意是排序后的第 k 大元素，不是第 k 个不同的元素。
请实现 KthLargest 类：
KthLargest(int k, int[] nums) 使用整数 k 和整数流 nums 初始化对象。
int add(int val) 将 val 插入数据流 nums 后，返回当前数据流中第 k 大的元素。
示例：
输入：
1
2
| ["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
|
输出：
1
2
3
4
5
6
7
8
9
| [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
|
堆~
简单数据结构，我们只需要维护一个大小为 k 的小根堆即可，每次 add 操作也只需要比较添加的元素和堆顶元素的大小即可，我写了 C++ 和 Python 的两个代码：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
| class KthLargest {
private:
priority_queue<int, vector<int>, greater<>> q;
int siz;
public:
KthLargest(int k, vector<int> &nums) {
siz = k;
for (auto i:nums) {
if (q.size() < k) q.emplace(i);
else if (i > q.top()) q.pop(), q.emplace(i);
}
}
int add(int val) {
if (q.size() < siz) {
q.push(val);
return q.top();
} else {
if (val < q.top()) return q.top();
else {
q.pop();
q.emplace(val);
return q.top();
}
}
}
};
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
| class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.siz = k
self.q = []
for i in nums:
if len(self.q) < self.siz:
heapq.heappush(self.q, i)
elif i > self.q[0]:
heapq.heappop(self.q)
heapq.heappush(self.q, i)
def add(self, val: int) -> int:
if len(self.q) < self.siz:
heapq.heappush(self.q, val)
return self.q[0]
else:
if val < self.q[0]:
return self.q[0]
else:
heapq.heappop(self.q)
heapq.heappush(self.q, val)
return self.q[0]
|