xiedaxia1hao
commented
2 years ago 模版题 295. Find Median from Data Stream
题目让我们建一个class去有效率的求一个data stream的中位数。该class里需要有两个函数：
void addNum(int num)：把num加到我们的class中double findMedian()：返回当前的data stream中的中位数 Example:
- insertNum(3)
- insertNum(1)
- findMedian() -> output: 2
- insertNum(5)
- findMedian() -> output: 3
- insertNum(4)
- findMedian() -> output: 3.5
这道题要怎么去思考呢？
最开始我们可以考虑用brute force的方法去做。我们可以维护一个sorted array，这样每次我们要找中文数的时候，往array中间的index找就可以了。但是这个方法我们需要O(n)的时间来让一个元素进入array。
我们能做的更好吗？因为我们实际上不需要保证整个array都是fully sorted的，我们只care中间的元素。
假设x是我们array中的中位数，那整个array中有一半的数字是小于（或等于）x的，另外一半的数字是大于（或等于）x的。那我们就可以考虑把array分成两个部分：第一个部分去保存小于x的numbers，另外一部分去保存大于x的numbers。那我们要返回的中位数，要么是第一部分数字中的最大值，或者是第二部分数字中的最小值，或者是两者的平均值。
所以我们很容易联想到用heap这个数据结构来保存我们的数字：
- 我们把第一部分小于x的数字保存到一个max heap中（方便我们取该部分的最大值）。
- 我们把第二部分大于x的数字保存到一个min heap中（方便我们取该部分的最小值）。
- 往heap中加新的元素需要O(lgn)，这个比brute force的方法更快。
- 而要返回中位数的时候，我们只要把两个heap中的peek取出并计算即可。
当data stream进来的number个数是偶数时，我们保证两个heap的size相同即可。如果是奇数的话，保证其中一个heap的size比另一个heap大1就好。这里选哪个heap都可以，我们这里assume我们的max heap的size可以比min heap的size大1。
我们可以走上面这个例子来看一下：
insertNum(3)：因为我们assume max heap的size可以比min heap的size大1，3我们直接入第一个max heap （如果incoming number小于等于max heap的peek，我们也需要加到我们的heap中）。要注意的是，每次加入一个新的元素，我们都需要把两个heap的size重新balance一下。就像上面说的，如果总number数是偶数，两个heap的size要一样，不然的话，max heap的size允许比min heap的元素大1（这样也方便我们直接取中位数，偶数的时候我们把两个heap的peek取平均就好，奇数的时候直接返回max heap的peek）。
insertNum(1): 因为1比3小，我们让1进max heap。
但是这时候两个heap的size不平衡了，我们要把max heap的最大值放到min heap中去。
findMedian(): 因为现在总numbers的数量是偶数，我们把max heap和min heap的peek值求个平均就好:(1+3)/2=2insertNum(5): 因为5比3大，我们把5加到min heap中，这时候min heap的size是2，max heap的size是1，不平衡了，我们需要重新balance一下，即把min heap的peek值放到max heap中。
findMedian(): 这个时候总的numbers的数量是奇数，我们直接返回max heap的peek值，即返回3.insertNum(4)：4比max heap的最大值要大，我们放入min heap中。此时两个heap的size是balanced，无需移动。
findMedian(): 这个时候总的numbers的数量是偶数，我们把max heap和min heap的peek值求个平均就好:(3+4)/2=3.5
代码如下：
class MedianFinder {
PriorityQueue<Integer> maxHeap; //containing first half of numbers
PriorityQueue<Integer> minHeap; //containing second half of numbers
public MedianFinder() {
// maxHeap = new PriorityQueue<>((a,b)->(b-a));
maxHeap = new PriorityQueue<>(Collections.reverseOrder());
minHeap = new PriorityQueue<>();
}
public void addNum(int num) {
if(maxHeap.isEmpty() || num <= maxHeap.peek()) {
maxHeap.offer(num);
} else {
minHeap.offer(num);
}
// either both the heaps will have equal number of elements or max-heap will have one
// more element than the min-heap
if(maxHeap.size() > minHeap.size()+1) {
minHeap.offer(maxHeap.poll());
} else if(maxHeap.size() < minHeap.size()) {
maxHeap.offer(minHeap.poll());
}
}
public double findMedian() {
if(maxHeap.size() == minHeap.size()) {
return (maxHeap.peek() + minHeap.peek())/2.0;
} else {
return maxHeap.peek();
}
}
}
对于有一些题目，我们会给到一堆可以分成两个parts的元素。去解决这类问题，我们通常对第一个part的最大值和第二个部分的最小值感兴趣。Two Heaps这个pattern很适合解决这类问题。
我们可以建一个max heap去保存第一个part的最大值，一个min heap去保存第二个part的最小值。
这么说有点抽象，我们可以通过下面这道经典例题来理解这个pattern。