每日算法：数据流的中位数

[[431427]]

本文转载自微信公众号「三分钟学前端」，作者sisterAn  。转载本文请联系三分钟学前端公众号。

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

例如，

[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 范围内，你将如何优化你的算法?

看到这个动态数组获取中位数问题，不要太激动，这太适合使用堆了，考察的就是堆的经典应用：中位数问题，详情可查看 前端进阶算法9：看完这篇，再也不怕堆排序、Top K、中位数问题面试了

解法：利用堆

解题思路：

这里需要维护两个堆：

• 大顶堆：用来存取前 n/2 个小元素，如果 n 为奇数，则用来存取前 Math.floor(n/2) + 1 个元素
• 小顶堆：用来存取后 n/2 个小元素

那么，根据题目要求，中位数就为：

• n 为奇数：中位数是大顶堆的堆顶元素
• n 为偶数：中位数是大顶堆的堆顶元素与小顶堆的堆顶元素的平均值

当数组为动态数组时，每当数组中插入一个元素时，都需要如何调整堆喃?

如果插入元素比大顶堆的堆顶要大，则将该元素插入到小顶堆中;如果要小，则插入到大顶堆中。

当插入完成后，如果大顶堆、小顶堆中元素的个数不满足我们已上的要求，我们就需要不断的将大顶堆的堆顶元素或小顶堆的堆顶元素移动到另一个堆中，直到满足要求

代码实现：

let MedianFinder = function() { 
    // 大顶堆，用来保存前 n/2 小的元素 
    this.lowHeap = new MaxHeap() 
    // 小顶堆，用来保存后 n/2 小的元素 
    this.hightHeap = new MinHeap() 
}; 
// 插入元素 
MedianFinder.prototype.addNum = function(num) { 
    // 如果大顶堆为空或大顶堆堆顶元素小于num，则插入大顶堆 
    // 否则插入到小顶堆中 
    if(!this.lowHeap.getSize() || num < this.lowHeap.getHead()) { 
        // 比大顶堆的堆顶小，插入到大顶堆中 
        this.lowHeap.insert(num) 
    } else { 
        // 比小顶堆的堆顶大，插入到小顶堆中 
        this.hightHeap.insert(num) 
    } 
 
    // 比较大小顶堆的是否依然保持平衡 
    if(this.lowHeap.getSize() - this.hightHeap.getSize() > 1) { 
        // 大顶堆往小顶堆迁移 
        this.hightHeap.insert(this.lowHeap.removeHead()) 
    } 
    if(this.hightHeap.getSize() > this.lowHeap.getSize()) { 
        // 小顶堆向大顶堆迁移 
        this.lowHeap.insert(this.hightHeap.removeHead()) 
    } 
}; 
// 获取中位数 
MedianFinder.prototype.findMedian = function() { 
    if(this.lowHeap.getSize() && this.lowHeap.getSize() === this.hightHeap.getSize()) { 
        return (this.lowHeap.getHead() + this.hightHeap.getHead())/2 
    } 
    return this.lowHeap.getHead() 
}; 

其中小顶堆定义：

// 小顶堆 
let MinHeap = function() { 
    let heap = [,] 
    // 堆中元素数量 
    this.getSize = ()=> heap.length - 1 
    // 插入 
    this.insert = (key) => { 
        heap.push(key) 
        // 获取存储位置 
        let i = heap.length-1 
        while (Math.floor(i/2) > 0 && heap[i] < heap[Math.floor(i/2)]) {   
            swap(heap, i, Math.floor(i/2)); // 交换  
            i = Math.floor(i/2);  
        } 
    } 
    // 删除堆头并返回 
    this.removeHead = () => { 
        if(heap.length > 1) { 
            if(heap.length === 2) return heap.pop() 
            let num = heap[1] 
            heap[1] = heap.pop() 
            heapify(1) 
            return num 
        } 
        return null 
    } 
    // 获取堆头 
    this.getHead = () => { 
        return heap.length > 1 ? heap[1]:null 
    } 
    // 堆化 
    let heapify = (i) => { 
        let k = heap.length-1 
        // 自上而下式堆化 
        while(true) { 
            let minIndex = i 
            if(2*i <= k && heap[2*i] < heap[i]) { 
                minIndex = 2*i 
            } 
            if(2*i+1 <= k && heap[2*i+1] < heap[minIndex]) { 
                minIndex = 2*i+1 
            } 
            if(minIndex !== i) { 
                swap(heap, i, minIndex) 
                i = minIndex 
            } else { 
                break 
            } 
        } 
    }  
    let swap = (arr, i, j) => { 
        let temp = arr[i] 
        arr[i] = arr[j] 
        arr[j] = temp 
    } 
} 

大顶堆定义：

// 大顶堆 
let MaxHeap = function() { 
    let heap = [,] 
    // 堆中元素数量 
    this.getSize = ()=>heap.length - 1 
    // 插入大顶堆 
    this.insert = (key) => { 
        heap.push(key) 
        // 获取存储位置 
        let i = heap.length-1 
        while (Math.floor(i/2) > 0 && heap[i] > heap[Math.floor(i/2)]) {   
            swap(heap, i, Math.floor(i/2)); // 交换  
            i = Math.floor(i/2);  
        } 
    } 
    // 获取堆头 
    this.getHead = () => { 
        return heap.length > 1 ? heap[1]:null 
    } 
    // 删除堆头并返回 
    this.removeHead = () => { 
        if(heap.length > 1) { 
            if(heap.length === 2) return heap.pop() 
            let num = heap[1] 
            heap[1] = heap.pop() 
            heapify(1) 
            return num 
        } 
        return null 
    } 
    // 堆化 
    let heapify = (i) => { 
        let k = heap.length-1 
        // 自上而下式堆化 
        while(true) { 
            let maxIndex = i 
            if(2*i <= k && heap[2*i] > heap[i]) { 
                maxIndex = 2*i 
            } 
            if(2*i+1 <= k && heap[2*i+1] > heap[maxIndex]) { 
                maxIndex = 2*i+1 
            } 
            if(maxIndex !== i) { 
                swap(heap, i, maxIndex) 
                i = maxIndex 
            } else { 
                break 
            } 
        } 
    }  
    let swap = (arr, i, j) => { 
        let temp = arr[i] 
        arr[i] = arr[j] 
        arr[j] = temp 
    } 
} 

复杂度分析：

时间复杂度：由于插入元素到堆的时间复杂度为 O(logn)，为树的高度;移动堆顶元素都需要堆化，时间复杂度也为O(logn);所以，插入( addNum )的时间复杂度为 O(logn) ，每次插入完成后求中位数仅仅需要返回堆顶元素即可， findMedian 时间复杂度为 O(1)

空间复杂度：O(n)

如果数据流中所有整数都在 0 到 100 范围内，我们可以尝试使用计数排序，但计数排序的时间复杂度是O(n + m)，其中 m 表示数据范围，复杂度较高，这里不适合，计数排序比较适合静态数组前k个最值问题 leetcode347：前 K 个高频元素

leetcode：https://leetcode-cn.com/problems/find-median-from-data-stream/solution/javascriptshu-ju-liu-de-zhong-wei-shu-by-user7746o/