【LeetCode 热题100道笔记】数据流的中位数

题目描述

中位数是有序整数列表中的中间值。如果列表的大小是偶数，则没有中间值，中位数是两个中间值的平均值。

例如 arr = [2,3,4] 的中位数是 3 。
例如 arr = [2,3] 的中位数是 (2 + 3) / 2 = 2.5 。
实现 MedianFinder 类:

MedianFinder() 初始化 MedianFinder 对象。

void addNum(int num) 将数据流中的整数 num 添加到数据结构中。

double findMedian() 返回到目前为止所有元素的中位数。与实际答案相差 10-5 以内的答案将被接受。

示例 1：
输入
[“MedianFinder”, “addNum”, “addNum”, “findMedian”, “addNum”, “findMedian”]
[[], [1], [2], [], [3], []]
输出
[null, null, null, 1.5, null, 2.0]

解释
MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1); // arr = [1]
medianFinder.addNum(2); // arr = [1, 2]
medianFinder.findMedian(); // 返回 1.5 ((1 + 2) / 2)
medianFinder.addNum(3); // arr[1, 2, 3]
medianFinder.findMedian(); // return 2.0

提示:

• $-10^5<=num<=10^5$
• 在调用 findMedian 之前，数据结构中至少有一个元素
• 最多 $5\ast 10^4$ 次调用 addNum 和 findMedian

思考

利用两个堆的特性拆分数据，实现中位数的快速定位：

• 小顶堆（minHeap）：存储数据流中较大的一半元素，堆顶是这部分的最小值（即所有元素的「上中位数候选」）。
• 大顶堆（maxHeap）：存储数据流中较小的一半元素，堆顶是这部分的最大值（即所有元素的「下中位数候选」）。
• 平衡规则：始终保持小顶堆的元素数量 ≥ 大顶堆，且两者数量差不超过 1。这样：
  • 当元素总数为奇数时，小顶堆堆顶就是中位数；
  • 当元素总数为偶数时，两堆顶的平均值就是中位数。

算法过程（以 addNum 和 findMedian 为核心）

1. 初始化（MedianFinder 构造函数）

• 创建自定义优先队列 MyPriorityQueue，分别初始化小顶堆（默认比较器 a - b，堆顶最小）和大顶堆（自定义比较器 b - a，堆顶最大）。
• 用 size 记录当前元素总数，辅助堆的平衡判断。

2. 插入元素（addNum 方法）

1. 确定插入堆：若元素 ≤ 小顶堆堆顶（或小顶堆为空），插入小顶堆（归为「较大一半」的最小端）；否则插入大顶堆（归为「较小一半」的最大端）。
2. 平衡堆大小：
   • 若小顶堆元素比大顶堆多 ≥ 2，将小顶堆堆顶移到大顶堆，保证差值 ≤ 1；
   • 若大顶堆元素比小顶堆多，将大顶堆堆顶移到小顶堆，保证小顶堆数量不小于大顶堆。

3. 查询中位数（findMedian 方法）

• 若元素总数为奇数（小顶堆 size > 大顶堆 size）：直接返回小顶堆堆顶；
• 若元素总数为偶数（两堆 size 相等）：返回两堆顶的平均值。

关键代码解析（自定义优先队列 MyPriorityQueue）

优先队列基于「堆」实现，核心是 swim（上浮，插入时维护堆序）和 sink（下沉，删除时维护堆序）操作：

方法	功能描述
front()	返回堆顶元素（不删除），即优先级最高的元素。
push(num)	插入元素：先加入堆尾，再通过 swim 上浮到正确位置，保证堆序。
pop()	删除堆顶：先交换堆顶与堆尾，删除堆尾，再通过 sink 下沉新堆顶到正确位置。
swim()	元素上浮：若当前元素优先级高于父节点，交换两者，直到满足堆序。
sink()	元素下沉：若当前元素优先级低于子节点，交换与优先级最高的子节点，直到满足堆序。

时空复杂度

操作	时间复杂度	空间复杂度
addNum	O(log n)	O(n)
findMedian	O(1)	O(n)

• 时间复杂度：堆的插入（push）和删除（pop）均为 O(log n)（n 为当前元素总数），findMedian 直接取堆顶，无需计算；
• 空间复杂度：两个堆共存储所有元素，总空间为 O(n)。

代码

var MedianFinder = function() {
    this.minHeap = new MyPriorityQueue();
    this.maxHeap = new MyPriorityQueue();
    this.maxHeap.setCompare((a, b) => b - a);
    this.size = 0;
};

/** 
 * @param {number} num
 * @return {void}
 */
MedianFinder.prototype.addNum = function(num) {
    this.size++;
    if (this.minHeap.isEmpty() || num <= this.minHeap.front()) {
        this.minHeap.push(num);
        if (this.maxHeap.size() + 1 < this.minHeap.size()) {
            this.maxHeap.push(this.minHeap.pop());
        }
    } else {
        this.maxHeap.push(num);
        if (this.maxHeap.size() > this.minHeap.size()) {
            this.minHeap.push(this.maxHeap.pop());
        }
    }
    
};

/**
 * @return {number}
 */
MedianFinder.prototype.findMedian = function() {
   // 当元素总数为奇数时，最大堆的堆顶就是中位数
    if (this.minHeap.size() > this.maxHeap.size()) {
        return this.minHeap.front();
    }
    // 当元素总数为偶数时，取两个堆顶的平均值
    else {
        return (this.maxHeap.front() + this.minHeap.front()) / 2;
    }
};

/** 
 * Your MedianFinder object will be instantiated and called as such:
 * var obj = new MedianFinder()
 * obj.addNum(num)
 * var param_2 = obj.findMedian()
 */

 class MyPriorityQueue {

    constructor(capacity = Number.MAX_SAFE_INTEGER, compare = (a, b) => a - b) {
        this._data = [];
        this._capacity = capacity;
        this._size = 0;
        this._compare = compare;
    }

    setCompare(compare) {
        this._compare = compare;
    }

    front() {
        return this._data[0];
    }

    push(num) {
        if (this._capacity === this._size) {
            this.pop();
        }
        this._data.push(num);
        this.swim();
        this._size++;
    }

    pop() {
        if (this._data.length === 0) return;
        [this._data[0], this._data[this._data.length-1]] = [this._data[this._data.length-1], this._data[0]];
        const item = this._data.pop();
        this.sink();
        this._size--;
        return item;
    }

    swim(index = this._data.length-1) {
        while (index > 0) {
            let pIndex = Math.floor((index-1)/2);
            if (this._compare(this._data[index],this._data[pIndex]) > 0) {
                [this._data[index], this._data[pIndex]] = [this._data[pIndex], this._data[index]];
                index = pIndex;
                continue;
            }
            break;
        }

    }

    sink(index = 0) {
        const n = this._data.length;
        while (true) {
            let left = 2 * index + 1;
            let right = left + 1;
            let biggest = index;

            if (left < n && this._compare(this._data[left], this._data[index]) > 0) {
                biggest = left;
            }
            if (right < n && this._compare(this._data[right], this._data[biggest]) > 0) {
                biggest = right;
            }

            if (biggest !== index) {
                [this._data[biggest], this._data[index]] = [this._data[index], this._data[biggest]];
                index = biggest;
                continue;
            }

            break;
        }
    }

    size() {
        return this._size;
    }

    isEmpty() {
        return this._size === 0;
    }

}