703. 数据流中的第 K 大元素
题目描述
设计类 KthLargest，它维护一个不断增长的整数集合，并能随时给出其中第 k 大的数。这里的第 k 大按把所有数从大到小排好后的第 k 个位置来算，相同的值各占一个位置。
KthLargest(k, nums)：用整数 k 和初始数组 nums 创建对象。add(val)：把 val 加入集合，然后返回当前集合中第 k 大的数。
测试数据以两个数组给出：方法名序列和对应参数；输出数组是每次调用的返回值，构造函数对应 null。
示例 1
输入：操作 = ["KthLargest","add","add","add","add","add"] 参数 = [[3,[4,5,8,2]],[3],[5],[10],[9],[4]] 输出：[null,4,5,5,8,8]
解释：初始为 [4,5,8,2]，k = 3；依次加入 3、5、10、9、4 后，第 3 大的数分别是 4、5、5、8、8。
示例 2
输入：操作 = ["KthLargest","add","add","add","add"], 参数 = [[4,[7,7,7,7,8,3]],[2],[10],[9],[9]] 输出：[null,7,7,7,8]
解释：k = 4，加入 2、10、9 后第 4 大仍为 7，再加入一个 9 后变为 8。
数据范围
0 <= nums.length <= 10^41 <= k <= nums.length + 1-10^4 <= nums[i] <= 10^4-10^4 <= val <= 10^4add最多调用10^4次。
解题思路
只保留最大的 k 个数，放在一个大小为 k 的最小堆里。堆顶是这 k 个数里最小的那个，正好就是全体的第 k 大。新数到来先压入堆，若堆超过 k 个，就弹出堆顶：被弹出的数至少比 k 个数小，以后再也不可能成为第 k 大。
为什么正确
每次 add 返回前，堆里恰好是目前最大的 k 个数（数量不足 k 时是全部），堆顶就是第 k 大。
复杂度分析
每次 add 时间 O(log k)，构造 O(n log k)；空间 O(k)。
易错点
要用最小堆而不是最大堆：我们关心的是“前 k 大里最小的那个”。另外初始 nums 可能少于 k 个，不能在构造时假设堆已满。
参考代码
Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.heap = []
for x in nums:
self.add(x)
def add(self, val: int) -> int:
heapq.heappush(self.heap, val)
if len(self.heap) > self.k:
heapq.heappop(self.heap)
return self.heap[0]class KthLargest {
private int k;
private PriorityQueue<Integer> heap;
public KthLargest(int k, int[] nums) {
this.k = k;
heap = new PriorityQueue<>();
for (int x : nums) {
add(x);
}
}
public int add(int val) {
heap.offer(val);
if (heap.size() > k) {
heap.poll();
}
return heap.peek();
}
}class KthLargest {
public:
int k;
priority_queue<int, vector<int>, greater<int>> heap;
KthLargest(int k, vector<int>& nums) {
this->k = k;
heap = priority_queue<int, vector<int>, greater<int>>();
for (int x : nums) {
add(x);
}
}
int add(int val) {
heap.push(val);
if ((int)heap.size() > k) {
heap.pop();
}
return heap.top();
}
};type minHeap []int
func (h minHeap) Len() int { return len(h) }
func (h minHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h minHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *minHeap) Push(x any) { *h = append(*h, x.(int)) }
func (h *minHeap) Pop() any { old := *h; x := old[len(old)-1]; *h = old[:len(old)-1]; return x }
type KthLargest struct {
k int
heap *minHeap
}
func Constructor(k int, nums []int) KthLargest {
this := KthLargest{k: k}
this.heap = &minHeap{}
for _, x := range nums {
this.Add(x)
}
return this
}
func (this *KthLargest) Add(val int) int {
heap.Push(this.heap, val)
if this.heap.Len() > this.k {
heap.Pop(this.heap)
}
return (*this.heap)[0]
}const heappush = (h, x) => {
h.push(x);
for (let i = h.length - 1, p; i > 0 && h[p = (i - 1) >> 1] > h[i]; i = p) [h[p], h[i]] = [h[i], h[p]];
};
const heappop = h => {
const top = h[0], last = h.pop();
if (h.length) {
h[0] = last;
for (let i = 0; ;) {
const l = 2 * i + 1, r = l + 1; let s = i;
if (l < h.length && h[l] < h[s]) s = l;
if (r < h.length && h[r] < h[s]) s = r;
if (s === i) break;
[h[s], h[i]] = [h[i], h[s]]; i = s;
}
}
return top;
};
var KthLargest = function(k, nums) {
this.k = k;
this.heap = [];
for (const x of nums) {
this.add(x);
}
};
KthLargest.prototype.add = function(val) {
heappush(this.heap, val);
if (this.heap.length > this.k) {
heappop(this.heap);
}
return this.heap[0];
};const heappush = (h: number[], x: number): void => {
h.push(x);
for (let i = h.length - 1, p: number; i > 0 && h[p = (i - 1) >> 1] > h[i]; i = p) [h[p], h[i]] = [h[i], h[p]];
};
const heappop = (h: number[]): number => {
const top = h[0], last = h.pop()!;
if (h.length) {
h[0] = last;
for (let i = 0; ;) {
const l = 2 * i + 1, r = l + 1; let s = i;
if (l < h.length && h[l] < h[s]) s = l;
if (r < h.length && h[r] < h[s]) s = r;
if (s === i) break;
[h[s], h[i]] = [h[i], h[s]]; i = s;
}
}
return top;
};
class KthLargest {
k: number;
heap: number[];
constructor(k: number, nums: number[]) {
this.k = k;
this.heap = [];
for (const x of nums) {
this.add(x);
}
}
add(val: number): number {
heappush(this.heap, val);
if (this.heap.length > this.k) {
heappop(this.heap);
}
return this.heap[0];
}
}