优先队列(一)数据流中的中位数
问题描述
设计一个支持如下两种操作的数据结构:
-
void addNum(int)
:从数据流中获取一个元素,添加到当前的数据结构中 -
double findMedian()
:返回当前数据结构中存储的数据的中位数
解决思路
由于这里无法确切知道数据元素的规模,因此一般通过列表的方式存储元素再求取中位数的方式不是特别可靠。
考虑使用 “堆” 数据结构来完成这个功能,维护两个堆:最大堆和最小堆,最大堆用于保存中位数元素以下的所有元素,最小堆保存大于等于中位数的所有元素。
为了方便,将最小堆中的堆顶元素视为奇数大小的数据流的中位数,在添加时注意小心地调整这两个堆中的元素以使得满足上面的条件
实现
class MedianFinder {
final PriorityQueue<Integer> minPq; // 维护最小堆,保存大于等于中位数的元素
final PriorityQueue<Integer> maxPq; // 维护最大堆,保存小于中位数的元素
public MedianFinder() {
minPq = new PriorityQueue<>((x, y) -> x - y);
maxPq = new PriorityQueue<>((x, y) -> y - x);
}
public void addNum(int num) {
/*
注意这里堆中元素的调整
*/
if (minPq.size() != maxPq.size()) {
minPq.offer(num);
maxPq.offer(minPq.poll());
} else {
maxPq.offer(num);
minPq.offer(maxPq.poll());
}
}
public double findMedian() {
/*
由于这两个堆满足了我们给定的条件,因此中位数的计算就变得简单
*/
if (minPq.size() != maxPq.size())
return minPq.peek();
return (minPq.peek() + maxPq.peek()) * 1.0 / 2;
}
}
复杂度分析:
-
时间复杂度:由于每次调用
addNum(int)
方法时都会触发堆的平衡操作,因此时间复杂度为 \(O(log_2n)\),对于中位数的计算,由于只是获取了两个堆的堆顶元素,因此时间复杂度为 \(O(1)\) -
空间复杂度:需要额外的空间来保存输入的数据,因此空间复杂度为 \(O(n)\)
参考:
[1] https://leetcode-cn.com/problems/find-median-from-data-stream/