~/DHRUVUpskilling
← board/DSA/Two Heap/dsa-two-heap-05
Solved·29 Aug

Find Median from a Data Stream

DifficultyMedium
PatternTwo Heap
TrackDSA
tl;dr

Create a data structure that can store a list of integers that can change in size over time and find the median from this dynamically growing list in constant time, O(1).

full write-up

Statement

Create a data structure that can store a list of integers that can change in size over time and find the median from this dynamically growing list in constant time, O(1).

Implement a class, MedianOfStream, which should support the following operations:

  1. Constructor() — This initializes the object of this class, which in turn creates the max heap and the min heap.
  2. InsertNum(num) — This adds an integer, num, to the data structure.
  3. FindMedian() — This finds the median of all elements seen so far. If there are an even number of elements, return the average of the two middle values.

Constraints

  • −10⁵ ≤ num ≤ 10⁵, where num is an integer received from the data stream.
  • There will be at least one element in the data structure before the median is computed.
  • At most 500 calls will be made to the function that calculates the median.

Examples

Example 1

Operations: InsertNum(5) FindMedian() → 5 InsertNum(3) FindMedian() → 4.0 InsertNum(8) FindMedian() → 5

Explanation:

  • After inserting 5: list = [5], median = 5.
  • After inserting 3: list = [3, 5], median = (3 + 5) / 2 = 4.0.
  • After inserting 8: list = [3, 5, 8], median = 5.

Example 2

Operations:

InsertNum(-1) InsertNum(2) FindMedian() → 0.5 InsertNum(-3) FindMedian() → -1

Explanation:

  • After inserting -1 and 2: list = [-1, 2], median = (-1 + 2) / 2 = 0.5.
  • After inserting -3: list = [-3, -1, 2], median = -1.

Find Median from a Number Stream

Idea Behind the Solution

Let's say x is the median of the numbers in an array.

  • Half of the numbers will be smaller than or equal to x.
  • The other half will be larger than or equal to x.

So, we can split the array into two halves:

  • One half stores the smaller numbers.
  • The other half stores the larger numbers.

The median will always be one of these two values:

  • The largest number in the small half, or
  • The smallest number in the large half.

If the total number of elements is even, the median is the average of these two numbers.

Why We Use Two Heaps

We need a data structure that can quickly tell us the smallest or largest number in a changing (growing) array. A heap is perfect for this.

So, we use two heaps:

  • A max heap for the smaller half of the numbers. This lets us quickly find the largest number among the smaller numbers.
  • A min heap for the larger half of the numbers. This lets us quickly find the smallest number among the larger numbers.

This is called the "two heaps" pattern.

How the Algorithm Works

  1. We store the first half of the numbers (the smaller ones) in a max heap.
  2. We store the second half of the numbers (the larger ones) in a min heap.
  3. To find the median, we just look at the top elements of both heaps. 4). We'll balance both heaps we want max size of maxheap to be minHeap +1 if it become invalid we'll move move top elements and move to other heap

A Small Trick for the Max Heap

In this code, we use a min heap twice. But we make one of them act like a max heap by storing all its numbers as negative values (we multiply each number by -1 before adding it).

This is a common trick, since many heap libraries only give a min heap by default.

Code

class medianOfStream {
  constructor() {
    this.maxHeapForSmallNum = new MinHeap();
    this.minHeapForLargeNum = new MinHeap();
  }

  insertNum(num) {
    if (this.maxHeapForSmallNum.size() == 0 || -1 * this.maxHeapForSmallNum.peek() >= num) {
      this.maxHeapForSmallNum.offer(-1 * num);
    } else {
      this.minHeapForLargeNum.offer(num);
    }

    if (this.maxHeapForSmallNum.size() > this.minHeapForLargeNum.size() + 1) {
      this.minHeapForLargeNum.offer(-1 * this.maxHeapForSmallNum.poll());
    } else if (this.maxHeapForSmallNum.size() < this.minHeapForLargeNum.size()) {
      this.maxHeapForSmallNum.offer(-1 * this.minHeapForLargeNum.poll());
    }
  }

  findMedian() {
    if (this.maxHeapForSmallNum.size() == this.minHeapForLargeNum.size()) {
      return (-1 * this.maxHeapForSmallNum.peek()) / 2.0 + this.minHeapForLargeNum.peek() / 2.0;
    }

    return -1 * (this.maxHeapForSmallNum.peek() / 1.0);
  }
}