[路飞]_leetcode_703. 数据流中的第 K 大元素

102 阅读1分钟

设计一个找到数据流中第 k 大元素的类(class)。注意是排序后的第 k 大元素,不是第 k 个不同的元素。

请实现 KthLargest 类:

KthLargest(int k, int[] nums) 使用整数 k 和整数流 nums 初始化对象。 int add(int val) 将 val 插入数据流 nums 后,返回当前数据流中第 k 大的元素。

解题思路

这一题可以通过小顶堆来维护数组的前k大数。

  1. 将元素插入到堆中
  2. 当堆长度大于 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;
        }
    }
}