[路飞]leetcode-703.数据流中的第 K 大元素

110 阅读1分钟

设计一个找到数据流中第 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

利用小顶堆的思路解题

class MinHeap {
  constructor(max) {
    this.arr = [];
    this.size = 0;
    this.max = max;
  }
  push(val) {
    if (this.size >= this.max && val < this.top()) return;
    this.arr.push(val);
    this.size++;
    if (this.size > 1) {
      let cur = this.size - 1;
      let parent = (cur - 1) >> 1;
      while (cur > 0 && this.arr[parent] > this.arr[cur]) {
        [this.arr[parent], this.arr[cur]] = [this.arr[cur], this.arr[parent]];
        cur = parent;
        parent = (cur - 1) >> 1;
      }
    }
    while (this.size > this.max) this.pop();
  }
  pop() {
    if (this.empty()) return;
    this.arr[0] = this.arr.pop();
    this.size--;
    let cur = 0,
      childl = cur * 2 + 1,
      childr = cur * 2 + 2;
    while (
      (childl < this.size && this.arr[childl] < this.arr[cur]) ||
      (childr < this.size && this.arr[childr] < this.arr[cur])
    ) {
      if (childr < this.size && this.arr[childr] < this.arr[childl]) {
        [this.arr[cur], this.arr[childr]] = [this.arr[childr], this.arr[cur]];
        cur = childr;
      } else {
        [this.arr[cur], this.arr[childl]] = [this.arr[childl], this.arr[cur]];
        cur = childl;
      }
      childl = cur * 2 + 1;
      childr = cur * 2 + 2;
    }
  }
  top() {
    return this.arr[0];
  }
  empty() {
    return this.size === 0;
  }
}


/**
 * @param {number} k
 * @param {number[]} nums
 */
 var KthLargest = function (k, nums) {
  this.heap = new MinHeap(k);
  for (let i = 0; i < nums.length; i++) {
    this.heap.push(nums[i]);
  }
};

KthLargest.prototype.add = function (val) {
  this.heap.push(val);
  return this.heap.top();
};