设计一个找到数据流中第 k 大元素的类(class)。注意是排序后的第 k 大元素,不是第 k 个不同的元素。
请实现 KthLargest 类:
KthLargest(int k, int[] nums) 使用整数 k 和整数流 nums 初始化对象。 int add(int val) 将 val 插入数据流 nums 后,返回当前数据流中第 k 大的元素。
解题思路
这一题可以通过小顶堆来维护数组的前k大数。
- 将元素插入到堆中
- 当堆长度大于 k 时,堆pop出最小值
代码
/**
* @param {number} k
* @param {number[]} nums
*/
var KthLargest = function(k, nums) {
this.h = new Heap();
this.k = k;
for (i of nums){
this.add(i);
}
};
/**
* @param {number} val
* @return {number}
*/
KthLargest.prototype.add = function(val) {
this.h.offer(val)
if (this.h.size() > this.k) this.h.pop();
return this.h.peek();
};
// 小顶堆
class Heap {
constructor() {
this.data = [];
}
swap(i, j) {
[this.data[i], this.data[j]] = [this.data[j], this.data[i]]
}
compartor(i, j) {
return this.data[i] - this.data[j];
}
size() {
return this.data.length;
}
peek() {
if (!this.size()) return null;
return this.data[0]
}
offer(val) {
this.data.push(val);
this.bubbleUp(this.data.length-1);
}
pop() {
if (!this.size()) return null;
let res = this.data[0];
if (this.size() > 1) {
this.data[0] = this.data.pop();
this.bubbleDown(0);
} else {
this.data.pop();
}
return res;
}
bubbleUp(index) {
while(index) {
let parent = (index - 1) >> 1;
if (this.compartor(index, parent) < 0) {
this.swap(index, parent);
index = parent;
} else {
break;
}
}
}
bubbleDown(index) {
let last = this.data.length - 1;
while(index < last) {
let temp = index, left = index * 2 + 1, right = index * 2 + 2;
if (left <= last && this.compartor(left, temp) < 0) {
temp = left;
}
if (right <= last && this.compartor(right, temp) < 0) {
temp = right;
}
if (temp == index) break;
this.swap(temp, index);
index = temp;
}
}
}