Find Median from Data Stream
Design a class MedianFinder. addNum(num) adds an integer from a stream, and findMedian() returns the median of every number added so far as a float. The median is the middle value of the numbers in sorted order; when the count is even, it is the average of the two middle values. Numbers can repeat and arrive in any order, and findMedian is only called after at least one addNum. Each example starts from a new, empty MedianFinder.
Examples
addNum(1); addNum(2); findMedian() returns 1.5
After adding 1 and 2 the sorted numbers are [1,2]. The count is even, so the median is (1 + 2) / 2 = 1.5.
addNum(1); addNum(2); findMedian(); addNum(3); findMedian() returns [1.5,2.0]
The output lists each findMedian result in call order. The first call sees [1,2] and returns 1.5; after 3 is added the numbers are [1,2,3], and the middle value 2 comes back as 2.0.
addNum(5); addNum(15); addNum(1); addNum(3); findMedian() returns 4.0
The numbers arrive out of order, but sorted they are [1,3,5,15]. The two middle values are 3 and 5, so the median is 4.0.
Constraints
- At most 5 * 10^4 calls to addNum and findMedian.
Hints
What if you could always look at the largest of the bottom half and the smallest of the top half in O(1)?
Maintain a max-heap for the lower half and a min-heap for the upper half, balanced in size.
On addNum: push to lower if it's <= lower's top, else upper. Rebalance so |lower| - |upper| stays in {0, 1}. findMedian: if lower has more elements, return its top; otherwise return the average of the two tops.
Drill This: Your Approach First
Approach
Keep the smaller half of the numbers in a max-heap (lower) and the larger half in a min-heap (upper), so the middle values sit at the two tops. addNum pushes to lower if the number is at most lower's top (or lower is empty), otherwise to upper, then moves one top across if needed so lower has the same count as upper or one more. findMedian returns lower's top when lower is bigger, else the average of both tops.
Time
O(log n) add, O(1) median
Space
O(n)
Pitfall
Solutions
class MedianFinder {
private PriorityQueue<Integer> lower = new PriorityQueue<>(Collections.reverseOrder());
private PriorityQueue<Integer> upper = new PriorityQueue<>();
public void addNum(int num) {
if (lower.isEmpty() || num <= lower.peek()) lower.offer(num);
else upper.offer(num);
if (lower.size() > upper.size() + 1) upper.offer(lower.poll());
else if (upper.size() > lower.size()) lower.offer(upper.poll());
}
public double findMedian() {
if (lower.size() > upper.size()) return lower.peek();
return (lower.peek() + upper.peek()) / 2.0;
}
}import heapq
class MedianFinder:
def __init__(self):
self.lower = [] # max-heap (negated)
self.upper = [] # min-heap
def addNum(self, num: int) -> None:
if not self.lower or num <= -self.lower[0]:
heapq.heappush(self.lower, -num)
else:
heapq.heappush(self.upper, num)
if len(self.lower) > len(self.upper) + 1:
heapq.heappush(self.upper, -heapq.heappop(self.lower))
elif len(self.upper) > len(self.lower):
heapq.heappush(self.lower, -heapq.heappop(self.upper))
def findMedian(self) -> float:
if len(self.lower) > len(self.upper):
return float(-self.lower[0])
return (-self.lower[0] + self.upper[0]) / 2.0