Efficient Solution to Sum of Squares Problem in TypeScript

A Step-by-Step Guide to Calculating the Running Median of a Number Stream with Optimized Algorithms

Problem

Compute the running median of a sequence of numbers. That is, given a stream of numbers, print out the median of the list so far on each new element.

Recall that the median of an even-numbered list is the average of the two middle numbers.

For example, given the sequence [2, 1, 5, 7, 2, 0, 5], your algorithm should print out:

2
1.5
2
3.5
2
2
2

Problem Breakdown

Median Definition:

For an odd-sized list, the median is the middle element.

For an even-sized list, the median is the average of the two middle elements.

Efficient Calculation:

We need to efficiently insert numbers into the list and then quickly find the median after each insertion.

Maintaining a sorted list would work, but re-sorting the list after each insertion is costly (O(n log n)).

A more efficient way to handle this is by using two heaps (min-heap and max-heap):

  • Max-heap stores the smaller half of the numbers.
  • Min-heap stores the larger half of the numbers.

Algorithm Explanation:

At each step, we:

  1. Insert the new number into one of the heaps.
  2. Ensure the heaps are balanced, i.e., the number of elements in the max-heap and min-heap differs by at most 1.
  3. Calculate the median:
  • If both heaps have the same size, the median is the average of the two roots (top elements) of the heaps.
  • If the max-heap has one more element than the min-heap, the median is the root of the max-heap.

Solution in TypeScript:

class RunningMedian {
  private minHeap: number[] = []; // larger half (min-heap)
  private maxHeap: number[] = []; // smaller half (max-heap, represented as a max-heap using negated values)

  constructor() {}

  // Helper to insert into heap (push and reheapify)
  private insertHeap(heap: number[], value: number, isMinHeap: boolean = true) {
    heap.push(value);
    this.heapifyUp(heap, heap.length - 1, isMinHeap);
  }

  // Helper to remove the root element from a heap (pop and reheapify)
  private removeHeapRoot(heap: number[], isMinHeap: boolean = true): number {
    const root = heap[0];
    const last = heap.pop();
    if (heap.length > 0 && last !== undefined) {
      heap[0] = last;
      this.heapifyDown(heap, 0, isMinHeap);
    }
    return root;
  }

  // Heapify up: restore the heap property after insertion
  private heapifyUp(heap: number[], index: number, isMinHeap: boolean) {
    while (index > 0) {
      const parentIndex = Math.floor((index - 1) / 2);
      if (isMinHeap ? heap[index] < heap[parentIndex] : heap[index] > heap[parentIndex]) {
        [heap[index], heap[parentIndex]] = [heap[parentIndex], heap[index]];
        index = parentIndex;
      } else {
        break;
      }
    }
  }

  // Heapify down: restore the heap property after removal
  private heapifyDown(heap: number[], index: number, isMinHeap: boolean) {
    const length = heap.length;
    while (index < length) {
      const leftChild = 2 * index + 1;
      const rightChild = 2 * index + 2;
      let target = index;

      if (leftChild < length && (isMinHeap ? heap[leftChild] < heap[target] : heap[leftChild] > heap[target])) {
        target = leftChild;
      }

      if (rightChild < length && (isMinHeap ? heap[rightChild] < heap[target] : heap[rightChild] > heap[target])) {
        target = rightChild;
      }

      if (target !== index) {
        [heap[index], heap[target]] = [heap[target], heap[index]];
        index = target;
      } else {
        break;
      }
    }
  }

  public addNumber(num: number): void {
    // Step 1: Add number to maxHeap or minHeap
    if (this.maxHeap.length === 0 || num <= -this.maxHeap[0]) {
      this.insertHeap(this.maxHeap, -num, false); // Negative to use max-heap property
    } else {
      this.insertHeap(this.minHeap, num, true); // Normal min-heap
    }

    // Step 2: Balance heaps (ensure maxHeap has at most 1 more element than minHeap)
    if (this.maxHeap.length > this.minHeap.length + 1) {
      const maxRoot = -this.removeHeapRoot(this.maxHeap, false);
      this.insertHeap(this.minHeap, maxRoot, true);
    } else if (this.minHeap.length > this.maxHeap.length) {
      const minRoot = this.removeHeapRoot(this.minHeap, true);
      this.insertHeap(this.maxHeap, -minRoot, false);
    }
  }

  public getMedian(): number {
    if (this.maxHeap.length > this.minHeap.length) {
      return -this.maxHeap[0];
    } else {
      return (-this.maxHeap[0] + this.minHeap[0]) / 2;
    }
  }
}

// Example usage:
const sequence = [2, 1, 5, 7, 2, 0, 5];
const runningMedian = new RunningMedian();

sequence.forEach(num => {
  runningMedian.addNumber(num);
  console.log(runningMedian.getMedian());
});

Explanation:

Heaps:

  • minHeap: A min-heap that holds the larger half of the numbers.
  • maxHeap: A max-heap (implemented using negated values) that holds the smaller half.

Balancing Heaps:

After inserting a number, we balance the heaps by making sure that the max-heap has at most one more element than the min-heap.

Median Calculation:

  • If the max-heap contains more elements, the median is the root of the max-heap.
  • If both heaps have an equal number of elements, the median is the average of the two roots.

Time Complexity:

  • Inserting into the heap takes O(log n) due to heap properties.
  • Finding the median is O(1) since we just need to access the root(s) of the heaps.
  • Thus, the overall time complexity per insertion is O(log n).

Alternative Solutions:

A brute force approach could sort the array after every insertion, but that would result in O(n log n) per insertion. The heap-based solution is more efficient for large input sizes.


If you like this, please think about buying me a coffee! ☕️.

Burhanuddin’s Code Compendium

Thank you for being a part of the Burhanuddin’s Code Compendium community! Before you go: