LeetCode 295. 数据流的中位数 堆题解，了解堆结构实例之一。
295. 数据流的中位数#
https://leetcode.cn/problems/find-median-from-data-stream/
中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。
例如，
[2,3,4] 的中位数是 3
[2,3] 的中位数是 (2 + 3) / 2 = 2.5
设计一个支持以下两种操作的数据结构：
void addNum(int num) - 从数据流中添加一个整数到数据结构中。
double findMedian() - 返回目前所有元素的中位数。
使用堆维护数据流#
之前在 堆结构实现及使用场景 中已经提到堆结构可以用来求分位数。
设我们的数据序列维护为递增的：
- 我们使用最大堆维护前半段小区间，堆顶此时为区间最大值；
- 使用最小堆维护后半段大区间，堆顶此时为区间最小值；
- 偶数个数时，两边维护元素数相同，奇数个数时，前半段维护多出来的一个元素，因此中位数为：
- 奇数个数：
maxHeap.peek()
- 偶数个数：
(maxHeap.peek+minHeap.peek) /2
核心问题：如何向两个堆中写入数据？
- 初始时，两个堆均为空，先向最大堆写入
- 下次进入时，再向最大堆写入，取堆顶，写入最小堆
- 再下次，两个堆
size一致，且都有数据，先向最小堆写入取出最小值，再向最大堆写入最小堆的堆顶数据
复杂度分析#
写：堆化复杂度O(logn)。
读：取堆顶复杂度O(1)。
使用了额外的堆结构存储数据，空间复杂度O(n)。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
|
public class MedianFinder {
/**
* 保存更小区间，这个区间存奇数个情况下多出来的数据
*/
PriorityQueue<Integer> maxHeap;
/**
* 保存更大区间
*/
PriorityQueue<Integer> minHeap;
public MedianFinder() {
maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
minHeap = new PriorityQueue<>(Comparator.naturalOrder());
}
public void addNum(int num) {
if (maxHeap.size() == minHeap.size()) {
// 等量的情况：
// 最大堆空时，先入最大堆
// 最大堆不空，先入最小堆，从最小堆中取出最小元素入最大堆
if (!maxHeap.isEmpty() && num > maxHeap.peek()) {
minHeap.offer(num);
maxHeap.offer(minHeap.poll());
} else {
maxHeap.offer(num);
}
} else {
// 最大堆中小元素个数大于最小堆中的大元素个数
// 这一步本质上是要往minHeap中补充数据，可以先入最大堆，从最大堆中取出最大元素入最小堆
maxHeap.offer(num);
minHeap.offer(maxHeap.poll());
}
}
public double findMedian() {
int maxInSmallerRange = maxHeap.peek() == null ? 0 : maxHeap.peek();
int minInBiggerRange = minHeap.peek() == null ? 0 : minHeap.peek();
if (maxHeap.size() == minHeap.size()) {
return (maxInSmallerRange + minInBiggerRange) / 2.0;
}
return maxInSmallerRange;
}
}
|
可以看到maxHeap.size() == minHeap.size()时，分之内代码可以简化为只保留23、24行，本质上是只往最大堆中写入元素，为了保证是整个序列的更小值，可以先往最小堆中写入，取最小堆堆顶即为整个序列的更小值，同时也是最大堆中的合法元素（num大于最大堆堆顶时，则插入了新的最大值，否则只是最大堆的中间元素）。