Heap¶
Table of Contents¶
295. Find Median from Data Stream¶
-
LeetCode | LeetCode CH (Hard)
-
Tags: two pointers, design, sorting, heap priority queue, data stream
295. Find Median from Data Stream - Python Solution
from heapq import heappop, heappush
# Dual Heaps
class MedianFinder:
def __init__(self):
self.minHeap = []
self.maxHeap = []
self.min_size = 0
self.max_size = 0
def addNum(self, num: int) -> None:
heappush(self.maxHeap, -num)
heappush(self.minHeap, -heappop(self.maxHeap))
self.min_size += 1
if self.min_size > self.max_size:
heappush(self.maxHeap, -heappop(self.minHeap))
self.min_size -= 1
self.max_size += 1
def findMedian(self) -> float:
if self.min_size == self.max_size:
return (-self.maxHeap[0] + self.minHeap[0]) / 2.0
return -self.maxHeap[0]
obj = MedianFinder()
obj.addNum(1)
obj.addNum(2)
assert obj.findMedian() == 1.5
obj.addNum(3)
assert obj.findMedian() == 2
obj.addNum(4)
assert obj.findMedian() == 2.5
obj.addNum(5)
assert obj.findMedian() == 3
print("All Passed.")
295. Find Median from Data Stream - C++ Solution
#include <iostream>
#include <queue>
#include <vector>
using namespace std;
class MedianFinder {
private:
priority_queue<int> maxHeap;
priority_queue<int, vector<int>, greater<int>> minHeap;
int max_size = 0;
int min_size = 0;
public:
MedianFinder() {}
void addNum(int num) {
if (min_size == max_size) {
minHeap.push(num);
maxHeap.push(minHeap.top());
minHeap.pop();
max_size++;
} else {
maxHeap.push(num);
minHeap.push(maxHeap.top());
maxHeap.pop();
min_size++;
}
}
double findMedian() {
if (min_size == max_size) {
return (maxHeap.top() + minHeap.top()) / 2.0;
} else {
return (double)maxHeap.top();
}
}
};
int main() {
MedianFinder* obj = new MedianFinder();
obj->addNum(1);
obj->addNum(2);
cout << obj->findMedian() << endl; // 1.5
obj->addNum(3);
cout << obj->findMedian() << endl; // 2
return 0;
}