educative.io

Median of a stream of numbers

I believe a simpler and cleaner implementation of the mxheap, minheap for median of a stream is as below:

from heapq import heappush, heappop
class MedianOfAStream:
  def __init__(self):
    self.xh, self.nh = [], [] #max, min heaps

  def insert_num(self, num):
    heappush(self.xh, -num)
    heappush(self.nh, -heappop(self.xh))

    if len(self.xh) < len(self.nh):
      heappush(self.xh, -heappop(self.nh))

  def find_median(self):
    if len(self.xh) > len(self.nh):
      return -self.xh[0]
    return (-self.xh[0] + self.nh[0])/2.

any idea why the solution in the lesson does minHeap.size() + 1)

The intent is to keep both the max and min heaps to be of the same size, at most off by a count of 1. So anytime the maxheap size is greater than the minheap size + 1, pop and re-adjust.
For this problem, I suggest you follow the diagrams shown in the solution carefully with the associated steps

2 Likes