To solve this problem, you need to find the median element many times.

To do this optimally, you should realize that each time you find the median, many of the numbers will be the same. For example, here are typical numbers that you might track over time:

# typical numbers over time
2
2 1
2 1 99
2 1 99 8
2 1 99 8 3

Since many of the same numbers appear each time, you should be able to optimize by remembering something about your past computations.

To do this, we should use a Heap. You can store two Heaps: one that has the bigger half of all the numbers, and one that has the smaller half.

This lets you quickly find the middle numbers - they're just the root nodes of the heaps.

We'll make sure the top heap always gets the extra element, if there are an odd number of elements. This means the median is the root node of the top heap if there are an odd number of elements. It also means the median is the average of the 2 root nodes if there are an even number of elements. Here's the code for this:

def getMedian(self):


    if isEven:
        return (topheap[0]+bottomheap[0])/2  # average of middle numbers

    if isOdd:
        return topheap[0]  # middle number
        # since topheap always gets the extra element

We have to code the add function. There are 2 cases we have to consider:

Case 1. You have an odd number of elements.

Open SpoilerClose Spoiler

In this case, the top heap will start with an extra element (that's one of our assumptions).

We need to add num to one of the heaps, and then rebalance the heaps so they're the same size again. Here are some examples you should consider to figure out how to do this:

if num == 7: # move 5 to bottom heap, and add 7 to top heap
if num == 6: # move 5 to bottom heap, and add 6 to top heap
if num == 5: # add 5 to bottom heap
if num == 4: # add 4 to bottom heap
if num == 3: # add 3 to bottom heap

We can generalize this with code, like this:

def add(num):
    if isOdd:
        if num <= topheap[0]:
            heapq.heappush(bottomheap, -num)
        else:
            heapq.heappush(bottomheap, -heapq.heappop(topheap))
            heapq.heappush(topheap, num)

# bottomheap is a maxheap
# to implement a Max-Heap, we just negate all the values and use a Min-Heap

Case 2. You have an even number of elements.

Open SpoilerClose Spoiler

In this case, both heaps will be the same size:

We need to make sure the top heap gets an extra number (by assumption). To do this, you need to consider examples:

if num == 6: # add 6 to top heap
if num == 5: # add 5 to top heap
if num == 4: # add 4 to top heap
if num == 3: # move 4 to top heap, and add 3 to bottom heap
if num == 2: # move 4 to top heap, and add 2 to bottom heap

Here's the code that generalizes this:

def add(num):
    if isEven:
        if num >= -bottomheap[0]:
            heapq.heappush(topheap, num)
        else:
            heapq.heappush(topheap, -heapq.heappop(bottomheap))
            heapq.heappush(bottomheap, -num)

And if we start out with no elements, we should add the first element we see to the top heap. Here's the full code:

from heapq import heappush, heappop

class MedianTracker:
    def __init__(self):
        # To implement a Max-Heap, we just negate all the values and use a Min-Heap, since Python only has Min-Heaps.
        self.topheap = []
        self.bottomheap = [] # max heap, contains negative values

    def add(self, num):
        topheap, bottomheap = self.topheap, self.bottomheap

        # check if empty and add to topheap
        if len(topheap) == 0:
            heappush(topheap,num)
            return

        isEven = len(topheap) == len(bottomheap)
        # Odd case
        if not isEven:
            if num <= topheap[0]:
                heappush(bottomheap, -num)
            else:
                heappush(bottomheap, -heappop(topheap))
                heappush(topheap, num)
        # Even case
        else:
            if num >= -bottomheap[0]:
                heappush(topheap, num)
            else:
                heappush(topheap, -heappop(bottomheap))
                heappush(bottomheap, -num)

    def getMedian(self):
        topheap, bottomheap = self.topheap, self.bottomheap
        isEven = len(topheap) == len(bottomheap)
        if isEven: return (-bottomheap[0] + topheap[0]) / 2
        else:      return topheap[0]

Add Function Time Complexity $O(\log n)$ for every number added. This is because the heaps contain $O(n/2) = O(n)$ numbers, which means they both have a height of O(log n). So adding/removing from the heaps is $O(\log n)$. When we add a new number, we add/remove from the heaps something like 3 times worst-case, so the total time complexity is $O(3\log n)$. But remember, big O notation ignores factors of 3, so this is the same as $O(\log n)$.

Find Median Time Complexity $O(1)$. The median comes directly from the 2 root nodes of the heaps, which we can access immediately.

Space Complexity $O(n)$ to store the heap.

Other solutions: