在上一篇《图解堆排序原理》中已经介绍过了堆排序的原理，那么这一篇就来看看怎么利用堆排序来解决数据流中位数问题吧。
由于中位数的计算，需要获取排序后数组的中间一位或者两位数，如果是通过维护一个有序数组：
- 每次添加新元素，都需要寻找合适的位置，将新元素插入有序数组中，可以使用二分法来找到插入位置，时间复杂度是
O(logn) - 找到插入位置后，还需要通过移动数组元素的方式，将新元素放置到数组中，平均时间复杂度为
O(n) - 最后可以计算出中位数的位置，计算出中位数的值。总体的一个时间复杂度还是
O(n)
那么使用两个数组呢？一个数组small按照升序存储数据流中较小的一半数据，另一个数组large按照降序存储数据流中较大的一半数据，设定数组small的元素数量len不少于large的元素数量。
- 当数据流的元素数量为奇数时，那么数据流的中位数便会是
small[len - 1] - 当数据流的元素数量为偶数时，那么数据流的中位数便会是
(samll[len - 1] + large[0]) / 2。
当插入新元素num的时候，为了保持small中的元素数量不小于large，需要分情况讨论:
-
如果两个数组长度相同，都为
len，应该给small数组添加元素- 如果数组长度为
0，那么直接将元素添加到small中 - 如果
num <= small[len-1]，那么num可以直接插入small数组中 - 如果
num > small[len-1]，说明num应该放入large数组中，可是这样的话large数组的元素数量会比small数组大1，需要将large数组中的最小的元素large[0]push到small数组中
- 如果数组长度为
-
如果两个数组长度不相同，则
small数组的长度为len，large数组的长度为len-1，应该给large数组添加元素- 如果
num >= small[len-1]，那么num可以直接插入large数组中 - 如果
num < small[len-1]，说明num应该放入small数组中，可是这样的话small数组的元素数量会比large数组大2，需要将small数组中的最大的元素small[len-1]push到large数组中。这样small和large数组的长度都为len了。 可以看到，即使使用了两个数组来保存数据流中较小和较大的一半，时间复杂度依然是O(n)，因为还是需要对元素进行排序和移动。
- 如果
那有什么办法可以优化吗？这就回到正题，堆！大顶堆的一个特性就是根元素为数组中最大的元素，并且所有的元素只需要大于等于子元素。也就是说，堆并不一定是有序的数组。但是求中位数最多只需要用到small中的最大元素及large数组中的最小元素，即大顶堆和小顶堆的根顶元素，时间复杂度为O(1)。插入新元素时，重新堆化，时间复杂度为O(logn)。
以数据流[3,4,7,1]为例
- 初始化大顶堆
small和小顶堆large。此时两个堆都没有元素，因此将3放入small中 - 插入元素
4时，small堆顶元素为3，4 > 3。因此4可以直接放入large中 - 插入元素
7时，7会大于small堆顶元素3，因此先将7放入large中，将large的堆顶元素弹出并且插入small中，将两个堆重新堆化，确保符合堆的性质。
4. 插入元素
1时，1小于small堆顶元素3，因此先将1放入small中，再将small的堆顶元素弹出并插入large中，将两个堆堆化。
在实际编码过程中，为了避免每次插入元素时，都去跟small或者large的堆顶元素做比较，可以直接插入要插入堆的相反堆，堆化后弹出堆顶元素，插入目标堆。
const defaultCmp = (x, y) => x > y; // 默认是最大堆
const swap = (arr, i, j) => [arr[i], arr[j]] = [arr[j], arr[i]];
// 堆结构
class Heap {
constructor(cmp){
this.container = [];
// 通过定义比较函数，可以分别创建大顶堆和小顶堆
this.cmp = cmp ?? defaultCmp;
}
// 插入元素
insert(data){
let {container, cmp} = this;
// 将元素push到数组末尾
container.push(data);
let index = container.length - 1;
// 自底向上调整涉及交换的节点
while(index){
let parent = Math.floor((index - 1)/2);
if(!cmp(container[parent], container[index])){
swap(container, parent, index);
index = parent;
}else{
// 如果父节点已经满足堆的性质了，调整结束
break;
}
}
}
// 弹出堆顶元素后，需要堆化，使剩余元素满足堆的性质
heapify(){
let {container, cmp} = this;
// 保存堆顶元素
let res = container[0];
// 将末尾元素与堆顶元素交换，避免移动所有元素
swap(container, 0, container.length-1);
// 弹出堆顶元素
container.pop();
// 自顶向下调整涉及交换的堆元素，exchange为左子节点
let index = 0, exchange = 2 * index + 1, len = container.length;
// 当涉及交换的元素在堆中，继续
while(exchange < len){
// 判断右子节点与左子节点，究竟哪个要参与交换
let right = 2 * index + 2;
if(right < len && cmp(container[right], container[exchange])){
exchange = right;
}
// 不符合堆性质时，父节点与子节点之一进行交换
if(!cmp(container[index],container[exchange])){
swap(container, index, exchange);
}else{
break;
}
// 继续判断被交换后的子节点是否满足堆性质
index = exchange;
exchange = 2 * index + 1;
}
return res;
}
// 获取堆顶元素
top(){
if(this.container.length) return this.container[0];
return null;
}
}
var MedianFinder = function() {
// 大顶堆
this.small = new Heap();
// 小顶堆
this.large = new Heap((x,y) => y > x);
};
MedianFinder.prototype.addNum = function(num) {
if(this.small.container.length !== this.large.container.length){
this.small.insert(num);
this.large.insert(this.small.top());
this.small.heapify();
}else{
this.large.insert(num);
this.small.insert(this.large.top());
this.large.heapify();
}
};
/**
* @return {number}
*/
MedianFinder.prototype.findMedian = function() {
let total = this.small.container.length + this.large.container.length;
if(total % 2 === 0){
return (this.small.top() + this.large.top()) / 2;
}
return this.small.top();
};
