Running Median Stream
Problem
Numbers arrive one at a time. After each arrival, report the median of everything received so far: the middle value when the count is odd, and the average of the two middle values when it is even. Return the list of medians, one per arrival, as decimal numbers. The values do not arrive in any particular order.
Examples
Input: stream = [5, 15, 1, 3]
Output: [5.0, 10.0, 5.0, 4.0]
Why: after three arrivals the values are 1, 5, 15 and the middle one is 5Input: stream = [1, 2]
Output: [1.0, 1.5]
Why: an even count averages the two middle valuesInput: stream = [7]
Output: [7.0]
Why: edge case, a single value is its own medianHints
0 / 3
Re-sorting after every arrival gives the right answer and is far too slow. Notice that the median only ever needs the values sitting next to the middle, not the full order.
Split the values into a lower half and an upper half. If you could always see the largest of the lower half and the smallest of the upper half, the median would be immediate.
Keep the lower half in a structure exposing its maximum and the upper half in one exposing its minimum. Route each arrival through both so it lands in the correct half, then move one value across if the halves get more than one apart. Read the median off the one or two exposed values.
Solution
Two heaps face each other across the middle: the lower half exposes its largest value and the upper half its smallest, so the median is always one or two values away. Pushing every arrival into the lower half and immediately handing its largest to the upper half guarantees the value lands on the correct side regardless of where it belongs, and one rebalancing move keeps the lower half never smaller than the upper. Each arrival costs a logarithm. Time is O(n log n) overall, and space is O(n).
import heapq
def running_medians(stream):
low, high = [], [] # low is a max-heap (negated), high is a min-heap
out = []
for x in stream:
heapq.heappush(low, -x) # everything enters through low
heapq.heappush(high, -heapq.heappop(low)) # hand low's largest to high
if len(high) > len(low): # keep low the bigger half
heapq.heappush(low, -heapq.heappop(high))
if len(low) > len(high):
out.append(float(-low[0]))
else:
out.append((-low[0] + high[0]) / 2)
return out
print(running_medians([5, 15, 1, 3])) # -> [5.0, 10.0, 5.0, 4.0]
print(running_medians([1, 2])) # -> [1.0, 1.5]
print(running_medians([7])) # -> [7.0]Stuck on the idea rather than the code? Priority Queue covers it.