Description
Support addNum and findMedian for an integer stream. Median queries occur only after insertion. Examples list operations and their outputs; insertion returns null.
Solution
from heapq import heappush, heappop
def create_median_finder():
return {'low': [], 'high': []}
def add_num(state, num):
heappush(state['low'],-num)
heappush(state['high'],-heappop(state['low']))
if len(state['high']) > len(state['low']):
heappush(state['low'],-heappop(state['high']))
def find_median(state):
if len(state['low']) > len(state['high']):
return -state['low'][0]
return (-state['low'][0]+state['high'][0])/2Examples
Inputs are positional arguments. Trees use level-order arrays; linked lists use value arrays. Design problems list operations in order.
Example 1
- Input
[["addNum",2],["addNum",8],["findMedian"]]- Output
[null,null,5]
The average of 2 and 8 is 5.
Example 2
- Input
[["addNum",-3],["findMedian"]]- Output
[null,-3]
The only stored value is the median.
Example 3
- Input
[["addNum",1],["addNum",3],["addNum",2],["findMedian"]]- Output
[null,null,null,2]
Sorting the three values places 2 in the middle.
Approach
Keep a max-heap for the lower half and min-heap for the upper half. The lower half has either equal size or one extra element.
Time & space
O(log n) per insertion; O(1) per median query; O(n) space for n stored values.