有一个源源不断地生成整数的数据流, 假设你有足够的空间来保存生成的数. 设计一个名叫 MedianHolder 的结构, 使之可以随时取得之前生成的所有数的中位数.
要求
- 如果
MedianHolder 已经保存了生成的 N 个数, 那么任意时刻将一个新数加入到 MedianHolder 的过程, 其时间复杂度是 O(logN)
- 取得已经产生的
N 个数整体的中位数的过程, 时间复杂度为 O(1)
题解
创建一个大根堆和一个小根堆, 尽量保证两个堆中元素数量均为 n/2 , 在添加元素时, 新元素若大于大根堆堆顶元素则放入小根堆, 若小于小根堆堆顶则放入大根堆, 这样即可保证前 n/2 个数都在大根堆中, 后 n/2 个数都在小根堆中, 中位数即为两堆的堆顶元素. 添加元素的同时确保小根堆的长度不小于大根堆, 若最后数据总数为奇数, 中位数即为小根堆堆顶元素; 若为偶数则为两堆堆顶元素和的一半.
添加新数的时间复杂度即为堆调整的时间复杂度, 即 O(logN) , 取得中位数的时间复杂度为 O(1)
Java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
| class MedianHolder {
private PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> b - a);
private PriorityQueue<Integer> minHeap = new PriorityQueue<>();
public void addNum(int num) {
minHeap.add(num);
maxHeap.add(minHeap.poll());
if (minHeap.size() < maxHeap.size())
minHeap.add(maxHeap.poll());
}
public double getMedian() {
if (maxHeap.size() == minHeap.size()) {
return (maxHeap.peek() + minHeap.peek()) / 2.0;
} else {
return minHeap.peek();
}
}
}
|
Python
heapq 中默认定义为小根堆, 通过保存相反数实现大根堆.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
| from heapq import *
class MedianFinder:
def __init__(self):
self.max_heap = []
self.min_heap = []
def addNum(self, num):
heappush(self.min_heap, num)
heappush(self.max_heap, -heappop(self.min_heap))
if len(self.min_heap) < len(self.max_heap):
heappush(self.min_heap, -heappop(self.max_heap))
def findMedian(self):
max_len = len(self.max_heap)
min_len = len(self.min_heap)
if min_len == 0:
return
return self.min_heap[0] if max_len != min_len else (-self.max_heap[0] + self.min_heap[0]) / 2.0
|