如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。
解题思路：
中位数又叫中值，是按顺序排列的一组数据中居于中间位置的数。
因为需要顺序排列，所以我们需要对数据流中的数字进行排序。
自定义规则
以中位数为界限，我们将小于中位数的数字成为较小数，大于中位数的数字成为较大数。较小数的那部分称为较小部分，较大数的部分称为较大部分。然后规定添加过程中，要保持较大部分的个数大于等于较小部分的个数。
我们在计算中位数的时候，需要分为两种情况，元素总个数为奇数和偶数：
- 元素总数为奇数时，需要拿到较大部分中的最小数作为中位数返回；
- 元素总数为偶数时，需要拿到较大部分的最小数和叫小部分的最大数的平均数返回；
数据结构选型
接下来就是思考如何存储较小部分，然后可以快速拿到较小部分的最大数，以及如何存储较大部分，然后可以快速拿到较大部分的最小数？
- 可以使用大根堆来存储较小部分，因为大根堆的根是所有堆内元素最大的数字；
- 可以使用小根堆来存储较大部分，因为小根堆的根是所有堆内元素最小的数字；
思路梳理
- 将数据流的数据，分为较大部分（用小根堆 $A$ 存储）和较小部分（用大根堆 $B$ 存储）。
- 添加数字 $x$ 的时候
- 如果此时总数为偶数，那么根据前面的规定，添加后 $A$ 要比 $B$ 多一个数字，先将 $x$ 添加到 $B$ 中，然后将 $B$ 的根添加到 $A$ 中。
- 如果此时总数为奇数，此时 $A.size > B.size$，那么为了维持 $A$ 和 $B$ 的大小平衡，添加后 $A$ 和 $B$ 的个数要想等，先将 $x$ 添加到 $A$ 中，然后将 $A$ 的根添加到 $B$ 中。
- 总数加一。
- 求中位数时
- 如果总数为奇数，那么直接返回 $A$ 的根。
- 如果总数为偶数，那么返回 $A$ 和 $B$ 的根的平均数。
代码：
class MedianFinder {
/**
* 小顶堆存放的是数据流排序后的后半部分，即较大的一半
*/
private final PriorityQueue<Integer> minHeap;
/**
* 大顶堆存放的是数据流排序后的前半部分，即较小的一半
*/
private final PriorityQueue<Integer> maxHeap;
private int total = 0;
/**
* initialize your data structure here.
*/
public MedianFinder() {
minHeap = new PriorityQueue<>(Comparator.naturalOrder());
maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
}
/**
* 添加的过程中，要保持小顶堆的元素个数大于等于大顶堆的个数，
* 即总元素个数为奇数的时候，小顶堆要比大顶堆多一个元素
*/
public void addNum(int num) {
// 如果是偶数，则入小顶堆，但是不能直接入
// 原因是，这个数字可能属于较小的一半，所以需要先入大顶堆，再把大顶堆的根弹出
if ((total & 1) == 0) {
maxHeap.offer(num);
minHeap.offer(maxHeap.poll());
} else {
// 如果是奇数，则入大顶堆，但是不能直接入
// 原因是，这个数字可能属于较大的一半，所以需要先入小顶堆，再把小顶堆的根弹出
minHeap.offer(num);
maxHeap.offer(minHeap.poll());
}
total++;
}
public double findMedian() {
// 偶数则返回中间两个元素的和的一半
if ((total & 1) == 0) {
return (double) (minHeap.peek() + maxHeap.peek()) / 2;
} else {
// 奇数则返回小顶堆的根
return (double) minHeap.peek();
}
}
}
复杂度分析
-
时间复杂度
- 添加数字 $O(logN)$：堆的弹出和压入使用 $O(logN)$ 的时间；
- 求中位数$O(1)$：获取堆顶元素使用 $O(1)$ 时间。
-
空间复杂度$O(N)$：其中 $N$ 为数据流中的元素数量，小顶堆 $A$ 和大顶堆 $B$ 最多同时保存 $N$ 个元素。
评论区