THN Interview Prep

295. Find Median from Data Stream

At a Glance

  • Topic: heap-priority-queue
  • Pattern: Top-K / Heap (two heaps)
  • Difficulty: Hard
  • Companies: Amazon, Google, Meta, Microsoft, Apple
  • Frequency: Very High
  • LeetCode: 295

Problem (one-liner)

Implement addNum(value) (streaming inserts) and findMedian() returning the median of all numbers seen so far in O(log n) amortized per insert and O(1) median query.

Recognition Cues

  • Streaming median / running median
  • Balance two heaps: max-heap holds smaller half, min-heap holds larger half
  • Invariant: heaps differ in size by at most one; median from tops

Diagram

At-a-glance flow (replace with problem-specific Mermaid as you refine this note). camelCase node IDs; no spaces in IDs.

Loading diagram…

Approaches

  • Brute force — buffer all values, sort on every findMedianO(n log n) per query.
  • Acceptable — balanced BST / order-statistic tree (policy-based or TreeMap multiset) — O(log n) insert + O(log n) median if supported.
  • Optimaltwo heaps: max-heap for smaller half, min-heap for larger half, sizes differ by at most one — O(log n) addNum, O(1) findMedian for standard heap APIs.

Optimal Solution

Go

import "container/heap"

type MinHeap []int

func (numbers MinHeap) Len() int { return len(numbers) }
func (numbers MinHeap) Less(indexA, indexB int) bool {
	return numbers[indexA] < numbers[indexB]
}
func (numbers MinHeap) Swap(indexA, indexB int) {
	numbers[indexA], numbers[indexB] = numbers[indexB], numbers[indexA]
}
func (numbers *MinHeap) Push(value any) { *numbers = append(*numbers, value.(int)) }
func (numbers *MinHeap) Pop() any {
	old := *numbers
	lastIndex := len(old) - 1
	item := old[lastIndex]
	*numbers = old[:lastIndex]
	return item
}

type MaxHeap []int

func (numbers MaxHeap) Len() int { return len(numbers) }
func (numbers MaxHeap) Less(indexA, indexB int) bool {
	return numbers[indexA] > numbers[indexB]
}
func (numbers MaxHeap) Swap(indexA, indexB int) {
	numbers[indexA], numbers[indexB] = numbers[indexB], numbers[indexA]
}
func (numbers *MaxHeap) Push(value any) { *numbers = append(*numbers, value.(int)) }
func (numbers *MaxHeap) Pop() any {
	old := *numbers
	lastIndex := len(old) - 1
	item := old[lastIndex]
	*numbers = old[:lastIndex]
	return item
}

type MedianFinder struct {
	lowerMax *MaxHeap
	upperMin *MinHeap
}

func Constructor() MedianFinder {
	lower := &MaxHeap{}
	upper := &MinHeap{}
	heap.Init(lower)
	heap.Init(upper)
	return MedianFinder{lowerMax: lower, upperMin: upper}
}

func (medianFinder *MedianFinder) AddNum(value int) {
	if medianFinder.lowerMax.Len() == 0 || value <= (*medianFinder.lowerMax)[0] {
		heap.Push(medianFinder.lowerMax, value)
	} else {
		heap.Push(medianFinder.upperMin, value)
	}
	if medianFinder.lowerMax.Len() > medianFinder.upperMin.Len()+1 {
		moved := heap.Pop(medianFinder.lowerMax).(int)
		heap.Push(medianFinder.upperMin, moved)
	} else if medianFinder.upperMin.Len() > medianFinder.lowerMax.Len()+1 {
		moved := heap.Pop(medianFinder.upperMin).(int)
		heap.Push(medianFinder.lowerMax, moved)
	}
}

func (medianFinder *MedianFinder) FindMedian() float64 {
	if medianFinder.lowerMax.Len() > medianFinder.upperMin.Len() {
		return float64((*medianFinder.lowerMax)[0])
	}
	if medianFinder.upperMin.Len() > medianFinder.lowerMax.Len() {
		return float64((*medianFinder.upperMin)[0])
	}
	return (float64((*medianFinder.lowerMax)[0]) + float64((*medianFinder.upperMin)[0])) / 2.0
}

JavaScript

class MinHeap {
	constructor() {
		this.values = [];
	}
	push(value) {
		this.values.push(value);
		this.bubbleUp(this.values.length - 1);
	}
	pop() {
		const top = this.values[0];
		const last = this.values.pop();
		if (this.values.length > 0) {
			this.values[0] = last;
			this.bubbleDown(0);
		}
		return top;
	}
	peek() {
		return this.values[0];
	}
	get size() {
		return this.values.length;
	}
	bubbleUp(index) {
		while (index > 0) {
			const parentIndex = Math.floor((index - 1) / 2);
			if (this.values[parentIndex] <= this.values[index]) {
				break;
			}
			[this.values[parentIndex], this.values[index]] = [
				this.values[index],
				this.values[parentIndex],
			];
			index = parentIndex;
		}
	}
	bubbleDown(index) {
		const length = this.values.length;
		while (true) {
			const leftIndex = index * 2 + 1;
			const rightIndex = index * 2 + 2;
			let smallest = index;
			if (
				leftIndex < length &&
				this.values[leftIndex] < this.values[smallest]
			) {
				smallest = leftIndex;
			}
			if (
				rightIndex < length &&
				this.values[rightIndex] < this.values[smallest]
			) {
				smallest = rightIndex;
			}
			if (smallest === index) {
				break;
			}
			[this.values[index], this.values[smallest]] = [
				this.values[smallest],
				this.values[index],
			];
			index = smallest;
		}
	}
}

class MaxHeap {
	constructor() {
		this.values = [];
	}
	push(value) {
		this.values.push(value);
		this.bubbleUp(this.values.length - 1);
	}
	pop() {
		const top = this.values[0];
		const last = this.values.pop();
		if (this.values.length > 0) {
			this.values[0] = last;
			this.bubbleDown(0);
		}
		return top;
	}
	peek() {
		return this.values[0];
	}
	get size() {
		return this.values.length;
	}
	bubbleUp(index) {
		while (index > 0) {
			const parentIndex = Math.floor((index - 1) / 2);
			if (this.values[parentIndex] >= this.values[index]) {
				break;
			}
			[this.values[parentIndex], this.values[index]] = [
				this.values[index],
				this.values[parentIndex],
			];
			index = parentIndex;
		}
	}
	bubbleDown(index) {
		const length = this.values.length;
		while (true) {
			const leftIndex = index * 2 + 1;
			const rightIndex = index * 2 + 2;
			let largest = index;
			if (
				leftIndex < length &&
				this.values[leftIndex] > this.values[largest]
			) {
				largest = leftIndex;
			}
			if (
				rightIndex < length &&
				this.values[rightIndex] > this.values[largest]
			) {
				largest = rightIndex;
			}
			if (largest === index) {
				break;
			}
			[this.values[index], this.values[largest]] = [
				this.values[largest],
				this.values[index],
			];
			index = largest;
		}
	}
}

class MedianFinder {
	constructor() {
		this.lowerMax = new MaxHeap();
		this.upperMin = new MinHeap();
	}
	addNum(value) {
		if (
			this.lowerMax.size === 0 ||
			value <= this.lowerMax.peek()
		) {
			this.lowerMax.push(value);
		} else {
			this.upperMin.push(value);
		}
		if (this.lowerMax.size > this.upperMin.size + 1) {
			this.upperMin.push(this.lowerMax.pop());
		} else if (this.upperMin.size > this.lowerMax.size + 1) {
			this.lowerMax.push(this.upperMin.pop());
		}
	}
	findMedian() {
		if (this.lowerMax.size > this.upperMin.size) {
			return this.lowerMax.peek();
		}
		if (this.upperMin.size > this.lowerMax.size) {
			return this.upperMin.peek();
		}
		return (this.lowerMax.peek() + this.upperMin.peek()) / 2;
	}
}

Walkthrough

Insert 1, 2. Lower holds [1], upper [2]. Median average.

Invariant: lowerMax size equals upperMin or is one larger (always push left-bias pattern above).

Edge Cases

  • First element
  • Duplicate values across halves — ties go either way if heaps rebalance
  • Integer overflow in median average — return float

Pitfalls

  • Wrong heap for which half — define clearly whether lower is max-heap (smaller numbers).
  • Forgetting heap.Init in Go before first use (constructor calls Init on empty heaps).
  • Off-by-one balance rule — allow size difference of at most one, not equal always.

Similar Problems

Variants

  • Sliding window median — evict old values (lazy deletion from heaps or multisets).
  • Weighted median — need different structure (sort cumulative weights).
  • Approximate quantiles — t-digest / sampling (systems track).

Mind-Map Tags

#two-heaps #median #streaming #heap-balance #order-statistics #hard-design

Last updated on

Spotted something unclear or wrong on this page?

On this page