[路飞]_数据流的中位数

388 阅读4分钟

「这是我参与2022首次更文挑战的第16天,活动详情查看:2022首次更文挑战

leetcode-295 数据流的中位数

题目介绍

中位数是有序列表中间的数。如果列表长度是偶数,中位数则是中间两个数的平均值。

例如,

[2,3,4] 的中位数是 3

[2,3] 的中位数是 (2 + 3) / 2 = 2.5

设计一个支持以下两种操作的数据结构:

  • void addNum(int num) - 从数据流中添加一个整数到数据结构中。
  • double findMedian() - 返回目前所有元素的中位数。

示例

addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3) 
findMedian() -> 2

进阶:

如果数据流中所有整数都在 0 到 100 范围内,你将如何优化你的算法?
如果数据流中 99% 的整数都在 0 到 100 范围内,你将如何优化你的算法?

解题思路

解法一:直接排序

解题代码

var MedianFinder = function() {
    this.arr = []
};

MedianFinder.prototype.addNum = function(num) {
    this.arr.push(num)
    this.arr.sort((a, b) => a - b)
};

MedianFinder.prototype.findMedian = function() {
    if (this.arr.length % 2 === 1) return this.arr[Math.floor(this.arr.length / 2)]
    return (this.arr[this.arr.length / 2] + this.arr[this.arr.length / 2 - 1]) / 2
};

直接排序在每次插入一个值之后,只需要将其插入到合适的位置即可,如果每次都要遍历整个数组进行排序的话,未免消耗太大,因此我们可以使用另一种方法,不需要每次都对所有的数进行排序

解法二:大顶堆 + 小顶堆

如果我们利用大顶堆保存前一半的数据,利用小顶堆保存后一半的数据,并且保持大顶堆和小顶堆的大小相等或者相差 1,那么中位数将在两个堆的堆顶元素中产生,要么是两个堆顶元素的和的一半,要么是其中一个堆顶元素,并且我们只需要在插入元素的过程中,始终维护两个堆的结构并且保证两个堆的大小相差始终小于等于 1 即可

解题步骤:

  1. 写一个 Heap 类,支持根据传入参数生成大顶堆或者小顶堆,堆中包含 size(堆的大小)、top(返回堆顶的值)、push(插入一个值)、pop(弹出堆顶的值)、sortBack(向上调整堆的结构)、sortFront(向下调整堆的结构)
  2. 创建一个大顶堆保存前半部分较小的值,创建一个小顶堆保存后半部分较大的值,并始终保持大顶堆的大小等于小顶堆的大小或者大顶堆的大小比小顶堆的大小小 1
  3. 如果新值比大顶堆的堆顶元素小,则将新值插入到大顶堆中,插入之后如果大顶堆的大小大于小顶堆的大小,则大顶堆不断弹出堆顶值插入到小顶堆中,直到大顶堆的大小不大于小顶堆的大小
  4. 如果新值比小顶堆的堆顶元素大,则将新值插入到小顶堆中,插入之后如果小顶堆的大小 - 大顶堆的大小 大于 1,则小顶堆不断弹出堆顶值插入到大顶堆中,直到小顶堆的大小 - 大顶堆的大小 小于等于 1
  5. 如果两个堆大小相等,则中位数为两个堆顶元素之和的一半;如果两个堆的大小不相等,则中位数为小顶堆的堆顶元素

数据流的中位数.gif

解题代码

var MedianFinder = function() {
    // 创建大顶堆用于保存前半部分较小的值
    this.frontHeap = new Heap(1)
    // 创建小顶堆用于保存后半部分较大的值
    this.backHeap = new Heap(-1)
};

MedianFinder.prototype.addNum = function(num) {
    // 需要先判断大顶堆不为空,才能判断其堆顶元素
    // 如果新值大于大顶堆的堆顶元素,则往大顶堆中插入
    if (this.frontHeap.size() && num < this.frontHeap.top()) {
        this.frontHeap.push(num)
        // 调整两个堆的大小关系
        while (this.frontHeap.size() > this.backHeap.size()) {
            this.backHeap.push(this.frontHeap.pop())
        }
    } else {
    // 如果大顶堆为空或者新值大于大顶堆的堆顶元素,则往小顶堆中插入
        this.backHeap.push(num)
        // 调整两个堆的大小关系
        while (this.backHeap.size() - this.frontHeap.size() > 1) {
            this.frontHeap.push(this.backHeap.pop())
        }
    }
};

MedianFinder.prototype.findMedian = function() {
    // 如果两个堆大小相等,则中位数为两个堆顶元素之和的平均数
    if (this.frontHeap.size() === this.backHeap.size()) return (this.frontHeap.top() + this.backHeap.top()) / 2
    // 如果两个堆大小不相等,则中位数为小顶堆的堆顶元素
    return this.backHeap.top()
};

class Heap {
    constructor(type) {
        this.arr = []
        // 根据 type 的值判断生成大顶堆还是生成小顶堆:-1 小顶堆  1 大顶堆
        this.type = type
    }

    // 返回堆的大小
    size() {
        return this.arr.length
    }

    // 返回堆顶元素
    top() {
        return this.arr[0]
    }

    // 往堆中插入元素
    push(val) {
        this.arr.push(val)
        this.sortBack()
    }

    // 弹出堆顶元素
    pop() {
        const val = this.arr[0]
        const back = this.arr.pop()
        if (this.size()) {
            this.arr[0] = back
            this.sortFront()
        }
        return val
    }

    // 向上调整堆结构
    sortBack() {
        let i = this.size() - 1
        if (this.type === -1) {
            while (i > 0 && this.arr[i] < this.arr[Math.floor((i - 1) / 2)]) {
                [this.arr[i], this.arr[Math.floor((i - 1) / 2)]] = [this.arr[Math.floor((i - 1) / 2)], this.arr[i]]
                i = Math.floor((i - 1) / 2)
            }
        } else {
            while (i > 0 && this.arr[i] > this.arr[Math.floor((i - 1) / 2)]) {
                [this.arr[i], this.arr[Math.floor((i - 1) / 2)]] = [this.arr[Math.floor((i - 1) / 2)], this.arr[i]]
                i = Math.floor((i - 1) / 2)
            }
        }
    }

    // 向下调整堆结构
    sortFront() {
        let i = 0
        if (this.type === -1) {
            while (i * 2 + 1 < this.size()) {
                let temp = i
                if (this.arr[temp] > this.arr[i * 2 + 1]) temp = i * 2 + 1
                if (i * 2 + 2 < this.size() && this.arr[temp] > this.arr[i * 2 + 2]) temp = i * 2 + 2
                if (temp === i) break
                [this.arr[temp], this.arr[i]] = [this.arr[i], this.arr[temp]]
                i = temp
            }
        } else {
            while (i * 2 + 1 < this.size()) {
                let temp = i
                if (this.arr[temp] < this.arr[i * 2 + 1]) temp = i * 2 + 1
                if (i * 2 + 2 < this.size() && this.arr[temp] < this.arr[i * 2 + 2]) temp = i * 2 + 2
                if (temp === i) break
                [this.arr[temp], this.arr[i]] = [this.arr[i], this.arr[temp]]
                i = temp
            }
        }
    }
}