Two Heaps — sliding median and k-th largest
Design a data structure that supports adding numbers and returning the median at any point.
The pattern
Split the stream into two halves:
max_pq— lower half. Python’sheapqis a min-heap, so negate values to simulate a max-heap.min_pq— upper half. Standard min-heap.
Invariant: every element in max_pq ≤ every element in min_pq. Sizes stay equal or max_pq has one extra.
Template
import heapq
min_pq = [] # upper half — standard min-heap
max_pq = [] # lower half — negate for max-heap
def add_to_min_pq(num):
"""Push num; if over capacity k, discard the smallest."""
if len(min_pq) < k:
heapq.heappush(min_pq, num)
else:
heapq.heappushpop(min_pq, num) # push then pop smallest — O(log k)
def add_to_max_pq(num):
"""Push num (negated); if over capacity k, discard the largest."""
if len(max_pq) < k:
heapq.heappush(max_pq, -num)
else:
heapq.heappushpop(max_pq, -num)
def add_to_min_heap(num):
"""Insert into min_pq; if num belongs in upper half, replace the top."""
if len(min_pq) < k:
heapq.heappush(min_pq, num)
elif num > min_pq[0]:
heapq.heapreplace(min_pq, num) # pop + push in one op, more efficient
def add_to_max_heap(num):
"""Insert into max_pq; if num is smaller than current max, replace."""
if len(max_pq) < k:
heapq.heappush(max_pq, -num)
return
top = -max_pq[0] # recover positive value
if top > num:
heapq.heapreplace(max_pq, -num)
def get_median():
if len(max_pq) == len(min_pq):
return (-max_pq[0] + min_pq[0]) / 2
return -max_pq[0] # max_pq has the extra element
Complexity
| Time | O(log n) per insertion, O(1) median query |
| Space | O(n) |
Key details
heappushpopvsheapreplace—heappushpoppushes first then pops (safe if heap is empty);heapreplacepops first then pushes (faster, but heap must be non-empty).- Rebalancing — after each insert, if sizes differ by more than 1, move the top of the larger heap to the smaller one.
- Negation — Python has no built-in max-heap. Negate on push, negate again on pop to recover the real value.
When this pattern shows up
- Find median from data stream (LeetCode 295)
- Sliding window median (LeetCode 480)
- IPO — maximize capital with k projects (LeetCode 502)
- Smallest range covering elements from k lists