Find Median from Data Stream
Support addNum and findMedian as numbers stream in, both fast.
Intuition
Keep the lower half in a max-heap and the upper half in a min-heap. The median is at the roots — no scanning. Each insert goes into one heap and rebalances at most one element across, so both roots always face the middle.
Two heaps is the standard answer whenever you need a running order statistic — a median, or the k-th smallest of a stream. The idea is to split the data at the answer: a max-heap holds the lower half so its root is the largest small value, a min-heap holds the upper half so its root is the smallest large value. The median is always sitting at one or both roots, so reading it is O(1).
Approach
Two heaps face each other
lo (max-heap) holds the smaller half, hi (min-heap) the larger. Invariant: every lo element ≤ every hi element, sizes differ by ≤ 1.
Insert then rebalance
Push into lo, move lo's max to hi (guaranteeing the ordering invariant), then if hi outgrew lo, move hi's min back. Two heap ops, done.
Median from the roots
Odd count → lo's root (lo is the bigger heap). Even → average of both roots. O(1).
Solution & live demo
Common pitfalls
Pushing straight onto the heap the value seems to belong to
if not self.lo or num <= -self.lo[0]:
heapq.heappush(self.lo, -num)
else:
heapq.heappush(self.hi, num)heapq.heappush(self.lo, -num) heapq.heappush(self.hi, -heapq.heappop(self.lo))
Both are correct, but the branchy version needs a separate rebalancing block and multiplies the cases you can get wrong. Pushing through lo and immediately relaying its largest element to hi enforces the ordering invariant unconditionally, with no comparison at all.
Forgetting to negate on the way out
return self.lo[0]
return -self.lo[0]
heapq is a min-heap only, so a max-heap is simulated by storing negatives. Every value crossing that boundary must flip sign — on the way in, on the way out, and on every transfer between the heaps.
Letting the heaps drift out of balance
heapq.heappush(self.lo, -num) heapq.heappush(self.hi, -heapq.heappop(self.lo))
heapq.heappush(self.lo, -num)
heapq.heappush(self.hi, -heapq.heappop(self.lo))
if len(self.hi) > len(self.lo):
heapq.heappush(self.lo, -heapq.heappop(self.hi))Without the third step every insert grows hi and shrinks lo, so the split stops falling at the middle and the roots stop being the median. Sizes must stay equal, or lo ahead by exactly one.
Edge cases
Rebalancing keeps halves even regardless of arrival order.
Invariant is ≤ so equal values may sit on both sides — fine.