Find Median from Data Stream hard

Problem Statement

Design a data structure that supports the following two operations:

• void addNum(int num): Add a number num to the data structure.
• double findMedian(): Return the median of all elements so far in the data structure.

Example 1:

MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1);    // arr = [1]
medianFinder.addNum(2);    // arr = [1, 2]
medianFinder.findMedian(); // return 1.5 (because median is (1 + 2) / 2 = 1.5)
medianFinder.addNum(3);    // arr = [1, 2, 3]
medianFinder.findMedian(); // return 2.0

Example 2:

MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1);    // arr = [1]
medianFinder.addNum(2);    // arr = [1,2]
medianFinder.addNum(3);    // arr = [1,2,3]
medianFinder.addNum(4);    // arr = [1,2,3,4]
medianFinder.findMedian(); // return 2.5

Steps

To efficiently find the median, we'll use two heaps: a max-heap (minHeap) to store the smaller half of the numbers, and a min-heap (maxHeap) to store the larger half. We maintain the invariant that minHeap.length <= maxHeap.length <= minHeap.length + 1. This ensures that the median is either the top of the maxHeap or the average of the tops of both heaps.

Explanation

1. addNum(num):

   • If maxHeap is empty or num is less than or equal to the top of maxHeap, add num to maxHeap. Otherwise, add num to minHeap.
   • After adding, balance the heaps: if maxHeap.length > minHeap.length + 1, move the largest element from maxHeap to minHeap. If minHeap.length > maxHeap.length, move the smallest element from minHeap to maxHeap.

2. findMedian():

   • If the total number of elements is odd, the median is the top of maxHeap.
   • If the total number of elements is even, the median is the average of the tops of maxHeap and minHeap.

Code

class MedianFinder {
    minHeap: number[] = [];
    maxHeap: number[] = [];

    addNum(num: number): void {
        if (this.maxHeap.length === 0 || num <= this.maxHeap[0]) {
            this.maxHeap.push(num);
            this.maxHeap.sort((a, b) => b - a); // Keep maxHeap sorted
        } else {
            this.minHeap.push(num);
            this.minHeap.sort((a, b) => a - b); //Keep minHeap sorted
        }

        //Balance heaps
        if (this.maxHeap.length > this.minHeap.length + 1) {
            this.minHeap.push(this.maxHeap.shift()!);
            this.minHeap.sort((a, b) => a - b);
        } else if (this.minHeap.length > this.maxHeap.length) {
            this.maxHeap.push(this.minHeap.shift()!);
            this.maxHeap.sort((a, b) => b - a);
        }
    }

    findMedian(): number {
        if (this.maxHeap.length === this.minHeap.length) {
            return (this.maxHeap[0] + this.minHeap[0]) / 2;
        } else {
            return this.maxHeap[0];
        }
    }
}


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

Complexity

• Time Complexity:

  • addNum(num): O(log n) due to heap operations (insertion and balancing).
  • findMedian(): O(1) - accessing the top of the heap is constant time.

• Space Complexity: O(n) to store the elements in the heaps. The heaps themselves take linear space in the worst case.