LeetCode Meditations: Find Median from Data Stream

Let's start with the description for this one:

The median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value, and the median is the mean of the two middle values.

• For example, for arr = [2, 3, 4], the median is 3.
• For example, for arr = [2, 3], the median is (2 + 3) / 2 = 2.5.

Implement the MedianFinder class:

• MedianFinder() initializes the MedianFinder object.
• void addNum(int num) adds the integer num from the data stream to the data structure.
• double findMedian() returns the median of all elements so far. Answers within 10^-5 of the actual answer will be accepted.

For example:

Input:
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]

Output:
[null, null, null, 1.5, null, 2.0]

Explanation:
MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1);    // arr = [1]
medianFinder.addNum(2);    // arr = [1, 2]
medianFinder.findMedian(); // return 1.5 (i.e., (1 + 2) / 2)
medianFinder.addNum(3);    // arr[1, 2, 3]
medianFinder.findMedian(); // return 2.0

This question is labeled as a hard one; however, finding the median itself is not hard.

The first idea that comes to mind is that we can continually add the numbers to an array, keep sorting it each time we do so, and return the median accordingly.

In fact, let's try it in TypeScript:

class MedianFinder {
  public nums: number[];

  constructor() {
    this.nums = [];
  }

  addNum(num: number): void {
    this.nums.push(num);
    this.nums.sort((a, b) => a - b);
  }

  findMedian(): number {
    let mid = Math.floor(this.nums.length / 2);

    if (this.nums.length % 2 === 0) {
      let mid1 = Math.floor(this.nums.length / 2) - 1;
      return (this.nums[mid] + this.nums[mid1]) / 2;
    } 

    return this.nums[mid];
  }
}

/**
 * Your MedianFinder object will be instantiated and called as such:
 * var obj = new MedianFinder()
 * obj.addNum(num)
 * var param_2 = obj.findMedian()
 */

Even though it passes some of the tests, it will end up with a Time Limit Exceeded error at one point, so we have to do better.

Indeed, let's take a deep breath, because we'll make use of heaps for an elegant solution.

When we think about it, if we have an even number of sorted numbers, the median will be the average of the maximum of the smaller half and the minimum of the larger half.

For example, if our numbers are these:

[3, 4, 7, 9]

the median is the average of 4 and 7 (which is 5.5).

The smaller half is [3, 4], while the larger half is [7, 9]. Indeed, 4 is the maximum in the smaller half and 7 is the minimum in the larger half.

However, if we have an odd length of numbers, we can just get either the maximum value from the smaller half, or the minimum value from the larger half.

Say, our numbers are these:

[3, 4, 7]

In this case, the smaller half can be [3, 4], and the larger half can be [7], so we can get the maximum value from the smaller half, which is 4, and be done with it.

Or, if the smaller half is just [3], and the larger half is [4, 7], we'll get the minimum value from the larger half, which is still 4.

You might already have an inkling that what we're talking about as the smaller half is a perfect candidate for a max heap because we're only concerned with the maximum value. Similarly, the larger half begs to be implemented as a min heap as we only deal with the minimum value in it. Of course, the elegance lies in the fact that we can get these values in constant time.

We are going to use TypeScript for the solution, and LeetCode's TypeScript environment (well, in fact, JavaScript environment) includes a handy package for using heaps (or, priority queues): Enter @datastructures-js/priority-queue!

We can use this package to create our max and min heaps, and make use of the functionality it gives such as enqueueing/dequeueing data.

So, inside the constructor of our MedianFinder class, we can initialize our min and max heaps:

class MedianFinder {
  public minHeap;
  public maxHeap;

  constructor() {
    this.minHeap = new MinPriorityQueue();
    this.maxHeap = new MaxPriorityQueue();
  }

  ...
}

Now, when it comes to adding a number, we need to choose a heap to add it into. We can choose either of them when we're starting off with empty heaps; however, let's use the max heap for adding a number first. That means, when we're going to add a number, we'll first try to put it in the max heap. But, we can only do it when the number we're going to add is less than the maximum value in the heap, or when the heap is empty. Otherwise, we can add it to the min heap.

So, we can have a handy getHeap function inside our class for choosing the heap to add the number into:

getHeap(n: number) {
  if (this.maxHeap.isEmpty() || n <= this.maxHeap.front().element) {
    return this.maxHeap;
  }

  return this.minHeap;
}

However, once we add the number, we need to make sure that the sizes of the heaps don't differ by more than 1. In that case, we'll dequeue the maximum value from the max heap, and enqueue it to the min heap:

if ((this.maxHeap.size() - this.minHeap.size()) > 1) {
  let num = this.maxHeap.dequeue().element;
  this.minHeap.enqueue(num);
}

If, however, the size of the min heap becomes larger than the size of the max heap, we need to rearrange things, and this time dequeue from the min heap to enqueue the value to the max heap:

if (this.maxHeap.size() < this.minHeap.size()) {
  let num = this.minHeap.dequeue().element;
  this.maxHeap.enqueue(num);
}

We can add these conditions to a separate rebalanceHeaps function for modularity's sake:

rebalanceHeaps(): void {
  if ((this.maxHeap.size() - this.minHeap.size()) > 1) {
    let num = this.maxHeap.dequeue().element;
    this.minHeap.enqueue(num);
  }

  if (this.maxHeap.size() < this.minHeap.size()) {
    let num = this.minHeap.dequeue().element;
    this.maxHeap.enqueue(num);
  }
}

And now, we're done with addNum itself. We only enqueued the new number to the proper heap, and rebalanced the heaps:

addNum(num: number): void {
  this.getHeap(num).enqueue(num);
  this.rebalanceHeaps();
}

Now, finding the median is a piece of cake. And, it's delicious because we'll get it in constant time:

findMedian(): number {
  let maxOfSmallerHalf = this.maxHeap.front().element;

  if (this.maxHeap.size() === this.minHeap.size()) {
    let minOfLargerHalf = this.minHeap.front().element;

    return (maxOfSmallerHalf + minOfLargerHalf) / 2;
  }

  return maxOfSmallerHalf;
}

And finally, here is the whole solution:

class MedianFinder {
  public minHeap;
  public maxHeap;

  constructor() {
    this.minHeap = new MinPriorityQueue();
    this.maxHeap = new MaxPriorityQueue();
  }

  addNum(num: number): void {
    this.getHeap(num).enqueue(num);
    this.rebalanceHeaps();
  }

  findMedian(): number {
    let maxOfSmallerHalf = this.maxHeap.front().element;

    if (this.maxHeap.size() === this.minHeap.size()) {
      let minOfLargerHalf = this.minHeap.front().element;

      return (maxOfSmallerHalf + minOfLargerHalf) / 2;
    }

    return maxOfSmallerHalf;
  }

  getHeap(n: number) {
    if (this.maxHeap.isEmpty() || n <= this.maxHeap.front().element) {
      return this.maxHeap;
    }

    return this.minHeap;
  }

  rebalanceHeaps(): void {
    if ((this.maxHeap.size() - this.minHeap.size()) > 1) {
      let num = this.maxHeap.dequeue().element;
      this.minHeap.enqueue(num);
    }

    if (this.maxHeap.size() < this.minHeap.size()) {
      let num = this.minHeap.dequeue().element;
      this.maxHeap.enqueue(num);
    }
  }
}

/**
 * Your MedianFinder object will be instantiated and called as such:
 * var obj = new MedianFinder()
 * obj.addNum(num)
 * var param_2 = obj.findMedian()
 */

This version is adapted from @kx0101's solution.

Time and space complexity

The time complexity is going to be $O(log \ n)$ because of the enqueue and dequeue operations on the heaps. However, getting the max value (or, the min value) is always $O(1)$ .

The space complexity is $O(n)$ as we need to create heaps to hold the numbers we have, and their sizes will grow as the length of the numbers increases.

And, this is the end of the chapter on heaps! Next up, we'll take a look at the technique of backtracking.
Until then, happy coding.