Find Median from Data Stream

The problem

Numbers arrive one at a time. Build a class MedianFinder with two methods:

• addNum(num) takes in the next number.
• findMedian() returns the median of all numbers so far.

The median is the middle value once the numbers are sorted. With an even count there are two middle values, and the median is their average.

Each example below is a list of calls on one new MedianFinder. Only findMedian returns something.

Example 1

Input: addNum(5), addNum(2), findMedian(), addNum(8), findMedian(), addNum(1), addNum(7), findMedian()

Output: 3.5, 5.0, 5.0

Sorted, the numbers are 2, 5, then 2, 5, 8, then 1, 2, 5, 7, 8. The picture shows the two halves at the end, with 5 on top of the smaller half.

	top
low	5	2	1
high	7	8

Example 2

Input: addNum(4), findMedian()

Output: 4.0

One number is its own median.

Example 3

Input: addNum(-1), addNum(-3), findMedian()

Output: -2.0

Two numbers: the average of −3 and −1 is −2.

Constraints

• -100,000 ≤ num ≤ 100,000
• findMedian is only called after at least one number was added.
• At most 50,000 calls in total.

Clarifying questions

Does the median have to be exact?

Close enough is fine: answers within 0.00001 are accepted. Averaging two ints as a float is exact enough.

Can numbers repeat?

Yes. Duplicates count separately, like in any sorted list.

Which call happens more?

Either can. So both should be fast, not one fast and one slow.

Can the numbers be negative?

Yes, down to −100,000. Nothing in the method depends on the sign.

Spot the pattern

The median only needs the middle of the sorted order, not the whole order. If the numbers are split into a smaller half and a larger half, the median is the largest of the small half, the smallest of the large half, or their average. “Keep the largest” and “keep the smallest” are what a heap does fast.

Solutions

Brute forceO(n) per addTwo heapsO(log n) per add

Show all

This problem has two approaches. Sorting on every findMedian is slower than the sorted list, so it is not a useful middle step.

Brute force O(n) per add

The idea

Keep every number in one sorted list. To add a number, find its place with binary search and insert it there. To find the median, look at the middle of the list.

The median is instant. But inserting into the middle of a list means every number after that spot must shift one place to make room. With n numbers already stored (n is the length of nums), that can be n shifts per add.

Walkthrough

nums (sorted)

answers

empty

1. nums = []
2. def addNum(num):
3. position = binary_search(nums, num)
4. nums.insert(position, num) # shifts what comes after
5. def findMedian():
6. nums_length, middle = len(nums), len(nums) // 2
7. if nums_length % 2: return nums[middle]
8. return (nums[middle - 1] + nums[middle]) / 2

Keep every number in one sorted list. Then the median is always in the middle.

← BackStep →Play1 / 10

Code

JSolution.javaPysolution.pyC#Solution.csGosolution.goTSsolution.ts

JavaPythonC#GoTypeScript

JavaPythonC#GoTypeScript

// Brute force: keep every number in one sorted list. The median is
// in the middle, but each insert shifts up to n numbers to make room.
class MedianFinder {
    private final List<Integer> nums = new ArrayList<>();

    public void addNum(int num) {
        // Binary search finds the spot fast. Making room is the slow part.
        int position = Collections.binarySearch(nums, num);
        nums.add(position < 0 ? -position - 1 : position, num);
    }

    public double findMedian() {
        int numsLength = nums.size(), middle = numsLength / 2;
        if (numsLength % 2 == 1) return nums.get(middle);
        return (nums.get(middle - 1) + nums.get(middle)) / 2.0;
    }
}

Big O

Time O(n) per addSpace O(n)Here is how we get these from the code, in four steps.

1. 1 Find the work that repeats

   The lit line inserts the number. Hidden inside it, every number after the insert spot moves one place to the right. Each move is a fixed job, so each shift counts as 1 step. The binary search before it takes only about log n steps, which is small next to the shifts.

   1. nums = []
   2. def addNum(num):
   3. position = binary_search(nums, num)
   4. nums.insert(position, num) # shifts what comes after
   5. def findMedian():
   6. nums_length, middle = len(nums), len(nums) // 2
   7. if nums_length % 2: return nums[middle]
   8. return (nums[middle - 1] + nums[middle]) / 2

2. 2 Count the runs on Example 1

   The table counts the shifts for each add in Example 1. Adding 1 shifts the most, because it goes in front of 2, 5 and 8. findMedian never shifts anything.

   0 + 1 + 0 + 3 + 1 = 5 shifts.

   addNum	5	2	8	1	7
   shifts	0	1	0	3	1

3. 3 Write the count for any n

   When k numbers are stored, a new number that is the smallest shifts all k of them. In the worst case, numbers arrive from largest to smallest, and the adds shift 0, 1, 2 and so on:

   one add: up to n shifts. n adds: 0 + 1 + … + (n − 1) = ½n² − ½n

   Check it with Example 1’s five numbers in the order 8, 7, 5, 2, 1: 0 + 1 + 2 + 3 + 4 = 10 = ½ × 25 − ½ × 5.

4. 4 Keep only the part that grows fastest

   For one add, the count is up to n, so O(n) per add. For a whole stream of n adds, drop the −½n and the ½ to get O(n²). findMedian is O(1).

   So the time is O(n) per add.

5. S Space: count the extra memory the same way

   The list holds every number, so the space is O(n).

   nums (n numbers)

Two heaps O(log n) per add

The idea

Keep two heaps. A heap is a tree kept in a list that always has its largest (or smallest) item on top, and it can add or remove an item in about log n steps.

• low is a max-heap with the smaller half of the numbers. Its top is the largest of them.
• high is a min-heap with the larger half. Its top is the smallest of them.

Two rules hold after every add. Every number in low is at most every number in high. And low has the same size as high, or one more.

To add a number, push it into low, then move the top of low into high. That sends the largest of the small numbers across, so the order rule holds. If high is now bigger, move its top back to low.

Then the median is on the tops: the top of low when the count is odd, or the average of both tops when it is even.

Why it is correct: by the two rules, low holds exactly the smaller half of the sorted order, so its top and the top of high are the one or two middle values.

Walkthrough

	top
low (max on top)
high (min on top)

answers

empty

1. low, high = max_heap(), min_heap()
2. def addNum(num):
3. low.push(num)
4. high.push(low.pop())
5. if len(high) > len(low):
6. low.push(high.pop())
7. def findMedian():
8. if len(low) > len(high): return low.top()
9. return (low.top() + high.top()) / 2

low keeps the smaller half with its largest on top. high keeps the larger half with its smallest on top. The middle of the data is always on the two tops.

← BackStep →Play1 / 13

Code

JSolution.javaPysolution.pyC#Solution.csGosolution.goTSsolution.ts

JavaPythonC#GoTypeScript

JavaPythonC#GoTypeScript

// Two heaps: low holds the smaller half (largest on top), high holds the
// larger half (smallest on top). The median sits on top. log n per add.
class MedianFinder {
    private final PriorityQueue<Integer> low = new PriorityQueue<>(Collections.reverseOrder());
    private final PriorityQueue<Integer> high = new PriorityQueue<>();

    public void addNum(int num) {
        // Pass the number through low, so the largest of the small half
        // moves up. This keeps every number in low <= every number in high.
        low.add(num);
        high.add(low.poll());
        // Keep the sizes equal, or low one bigger.
        if (high.size() > low.size()) low.add(high.poll());
    }

    public double findMedian() {
        if (low.size() > high.size()) return low.peek();
        return (low.peek() + high.peek()) / 2.0;
    }
}

Big O

Time O(log n) per addSpace O(n)Here is how we get these from the code, in four steps.

1. 1 Find the work that repeats

   The lit lines push into a heap or pop from it. A heap keeps its order by moving the new or removed item up or down its tree, one level per step, and the tree has about log n levels. So each push or pop costs up to log n steps. We count pushes and pops.

   1. low, high = max_heap(), min_heap()
   2. def addNum(num):
   3. low.push(num)
   4. high.push(low.pop())
   5. if len(high) > len(low):
   6. low.push(high.pop())
   7. def findMedian():
   8. if len(low) > len(high): return low.top()
   9. return (low.top() + high.top()) / 2

2. 2 Count the runs on Example 1

   The table counts pushes and pops per add on Example 1. An add always does a push, a pop and a push. When high gets too big, it does one more pop and push.

   5 + 3 + 5 + 3 + 5 = 21 pushes and pops.

   addNum	5	2	8	1	7
   heap pushes + pops	5	3	5	3	5

3. 3 Write the count for any n

   Each add does at most 5 pushes and pops, and each one costs at most about log n steps, where n is the count so far:

   one add: at most 5 × log n steps. findMedian: 1 step

   Check it with Example 1: no add did more than 5 pushes and pops, and the heaps never held more than 3 numbers, so each push or pop moved at most 1 level.

4. 4 Keep only the part that grows fastest

   Drop the 5. A fixed number of heap moves does not change the growth. What is left is O(log n) per add, and O(1) for findMedian. A stream of n adds costs O(n log n), against O(n²) for the sorted list.

   So the time is O(log n) per add.

5. S Space: count the extra memory the same way

   Every number sits in exactly one of the two heaps, so together they hold n numbers. The space is O(n).

   low (about n / 2)high (about n / 2)

Common mistake

Dividing two integers

In Java, C# and Go, (2 + 5) / 2 is integer division and gives 3, not 3.5. Divide by 2.0, or convert to a floating-point number first. In Python, use /, not //.

Tested at the limits

2,000 random call sequences of up to 21 calls, with numbers from -3 to 3, so equal numbers and medians like 0.5 come up. The brute force and the two heaps gave the same medians every time, and they matched a simple check that sorts all the numbers.

Largest inputs	Result
50,000 calls (the most allowed): 40,000 adds, many at the very ends of the range, and a median after every fourth add. Two heaps.	0.03 s, and about 2 MB of extra memory: the two heaps hold every number, which is the O(n) space.
50,000 calls: 40,000 rising numbers, so every add moves a number between the heaps, then 10,000 medians. Two heaps.	0.03 s.
50,000 calls, the first mix. Brute force.	0.07 s, about twice as slow. Each add shifts up to 40,000 numbers, but Python does that shift as one fast memory copy, so it stays quick at this size.