Median of Data Stream Using Two Heaps in DSA Go

Mental Model

We keep two groups of numbers: smaller half and larger half. The middle value is easy to find by looking at the tops of these groups.

Analogy: Imagine two boxes: one holds all smaller toys, the other holds bigger toys. The median toy is the one right between these two boxes.

MaxHeap (smaller half)      MinHeap (larger half)
      [5, 3, 1]                  [6, 8, 9]
       ↑ top is 5               ↑ top is 6

Median is either top of MaxHeap or average of tops.

Dry Run Walkthrough

Input: Stream: 5, 3, 8, 9, 1, 6

Goal: Find median after each new number is added

Step 1: Add 5 to MaxHeap (smaller half)

MaxHeap: [5] ↑top=5  MinHeap: [] null

Why: Start by adding first number to smaller half

Step 2: Add 3 to MaxHeap, then move top to MinHeap to balance

MaxHeap: [3] ↑top=3  MinHeap: [5] ↑top=5

Why: Keep heaps balanced so sizes differ by at most 1

Step 3: Add 8 to MinHeap, then move top to MaxHeap to balance

MaxHeap: [5, 3] ↑top=5  MinHeap: [8] ↑top=8

Why: Balance heaps after adding to larger half

Step 4: Add 9 to MinHeap

MaxHeap: [5, 3] ↑top=5  MinHeap: [8, 9] ↑top=8

Why: Add to larger half since 9 > 5, now sizes equal (2-2)

Step 5: Add 1 to MaxHeap

MaxHeap: [5, 3, 1] ↑top=5  MinHeap: [8, 9] ↑top=8

Why: Add to smaller half since 1 <= 5, sizes now 3-2 (differ by 1)

Step 6: Add 6 to MinHeap

MaxHeap: [5, 3, 1] ↑top=5  MinHeap: [6, 8, 9] ↑top=6

Why: Add to larger half since 6 > 5, sizes now equal (3-3), no rebalance needed

Result:

MaxHeap: [5, 3, 1] ↑top=5  MinHeap: [6, 8, 9] ↑top=6
Median after each insertion: 5, 4, 5, 6.5, 5, 5.5

Annotated Code

DSA Go

package main

import (
	"container/heap"
	"fmt"
)

// IntHeap is a min-heap of ints.
type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
	n := len(*h)
	x := (*h)[n-1]
	*h = (*h)[:n-1]
	return x
}

// MaxHeap is a max-heap of ints by inverting Less
type MaxHeap []int

func (h MaxHeap) Len() int           { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h MaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *MaxHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h *MaxHeap) Pop() interface{} {
	n := len(*h)
	x := (*h)[n-1]
	*h = (*h)[:n-1]
	return x
}

// MedianFinder holds two heaps
// maxHeap for smaller half, minHeap for larger half
// sizes differ by at most 1

type MedianFinder struct {
	maxHeap *MaxHeap
	minHeap *IntHeap
}

func Constructor() MedianFinder {
	maxH := &MaxHeap{}
	minH := &IntHeap{}
	heap.Init(maxH)
	heap.Init(minH)
	return MedianFinder{maxHeap: maxH, minHeap: minH}
}

func (mf *MedianFinder) AddNum(num int) {
	if mf.maxHeap.Len() == 0 || num <= (*mf.maxHeap)[0] {
		heap.Push(mf.maxHeap, num) // add to smaller half
	} else {
		heap.Push(mf.minHeap, num) // add to larger half
	}

	// Balance sizes
	if mf.maxHeap.Len() > mf.minHeap.Len()+1 {
		val := heap.Pop(mf.maxHeap).(int) // move top from maxHeap to minHeap
		heap.Push(mf.minHeap, val)
	} else if mf.minHeap.Len() > mf.maxHeap.Len() {
		val := heap.Pop(mf.minHeap).(int) // move top from minHeap to maxHeap
		heap.Push(mf.maxHeap, val)
	}
}

func (mf *MedianFinder) FindMedian() float64 {
	if mf.maxHeap.Len() == 0 {
		return 0.0
	}
	if mf.maxHeap.Len() > mf.minHeap.Len() {
		return float64((*mf.maxHeap)[0])
	} else {
		return (float64((*mf.maxHeap)[0]) + float64((*mf.minHeap)[0])) / 2.0
	}
}

func main() {
	mf := Constructor()
	nums := []int{5, 3, 8, 9, 1, 6}
	for _, num := range nums {
		mf.AddNum(num)
		median := mf.FindMedian()
		fmt.Printf("Added %d, current median: %.1f\n", num, median)
	}
}

if mf.maxHeap.Len() == 0 || num <= (*mf.maxHeap)[0] {

decide which heap to add new number to based on comparison with maxHeap top

heap.Push(mf.maxHeap, num)

add number to smaller half maxHeap

heap.Push(mf.minHeap, num)

add number to larger half minHeap

if mf.maxHeap.Len() > mf.minHeap.Len()+1 {

check if smaller half is too big and rebalance

val := heap.Pop(mf.maxHeap).(int)

remove top from smaller half to move to larger half

heap.Push(mf.minHeap, val)

add removed top to larger half

else if mf.minHeap.Len() > mf.maxHeap.Len() {

check if larger half is too big and rebalance

val := heap.Pop(mf.minHeap).(int)

remove top from larger half to move to smaller half

heap.Push(mf.maxHeap, val)

add removed top to smaller half

if mf.maxHeap.Len() > mf.minHeap.Len() {

if smaller half bigger, median is top of maxHeap

return float64((*mf.maxHeap)[0])

return median as top of smaller half

return (float64((*mf.maxHeap)[0]) + float64((*mf.minHeap)[0])) / 2.0

if equal size, median is average of tops

Output

Added 5, current median: 5.0

Added 3, current median: 4.0

Added 8, current median: 5.0

Added 9, current median: 6.5

Added 1, current median: 5.0

Added 6, current median: 5.5

Complexity Analysis

Time: O(log n) per insertion because heap push/pop operations take logarithmic time

Space: O(n) because all numbers are stored in the two heaps

vs Alternative: Compared to sorting all numbers each time (O(n log n)), using two heaps is faster for streaming data

Edge Cases

empty stream

FindMedian returns 0 if called before any number is added

DSA Go

if mf.maxHeap.Len() == 0 { return 0.0 }

all numbers equal

Heaps balance normally, median is that number

DSA Go

if mf.maxHeap.Len() > mf.minHeap.Len()+1 {

stream with one number

Median is that single number

DSA Go

if mf.maxHeap.Len() > mf.minHeap.Len() {

When to Use This Pattern

When you need to find the median of a growing list efficiently, use two heaps to keep track of smaller and larger halves for quick median retrieval.

Common Mistakes

Mistake: Not balancing the heaps after each insertion, causing incorrect median calculation
Fix: Always rebalance heaps so their sizes differ by at most one after each insertion

Mistake: Adding all numbers to one heap without comparison
Fix: Compare new number with maxHeap top to decide which heap to add to

Summary

Keeps track of median in a data stream using two heaps for smaller and larger halves.

Use when you want to find median quickly after each new number in a stream.

The median is always at the top of one or both heaps because they split the data into two balanced halves.

Practice

1. What is the main purpose of using two heaps (a max-heap and a min-heap) in finding the median of a data stream?

easy

A. To store all numbers in a single data structure

B. To sort all numbers in ascending order

C. To efficiently keep track of the lower half and upper half of the numbers separately

D. To find the maximum number in the data stream

Solution

1. Step 1: Understand the role of two heaps

   The max-heap stores the smaller half of the numbers, and the min-heap stores the larger half.

2. Step 2: Explain how this helps find the median

   By balancing these heaps, the median is either the top of one heap or the average of the tops of both heaps.

3. Final Answer:

   To efficiently keep track of the lower half and upper half of the numbers separately -> Option C

4. Quick Check:

   Two heaps split data halves = A [OK]

Hint: Two heaps split data halves for median [OK]

Common Mistakes:

• Thinking heaps sort all numbers fully
• Using one heap only
• Confusing max-heap and min-heap roles

2. Which of the following Go code snippets correctly initializes a max-heap and a min-heap for integers using container/heap package?

easy

A. maxHeap := &IntMaxHeap{}; minHeap := &IntMinHeap{}

B. maxHeap := heap.New(IntMaxHeap{}); minHeap := heap.New(IntMinHeap{})

C. maxHeap := &IntMinHeap{}; minHeap := &IntMaxHeap{}

D. maxHeap := make([]int, 0); minHeap := make([]int, 0)

Solution

1. Step 1: Recognize heap initialization in Go

   Heaps are initialized as pointers to custom heap types that implement heap.Interface.

2. Step 2: Identify correct syntax

   Using &IntMaxHeap{} and &IntMinHeap{} creates pointers to heap structs ready for heap operations.

3. Final Answer:

   maxHeap := &IntMaxHeap{}; minHeap := &IntMinHeap{} -> Option A

4. Quick Check:

   Pointer to heap struct = A [OK]

Hint: Use pointers to heap structs for container/heap [OK]

Common Mistakes:

• Using heap.New which does not exist
• Swapping max-heap and min-heap types
• Initializing as slices without heap.Interface

3. Given the following sequence of numbers inserted into the data stream: 5, 15, 1, 3, what is the median after all insertions using two heaps?

medium

A. 5

B. 8

C. 3

D. 4

Solution

1. Step 1: Insert numbers and balance heaps

   Insert 5 (maxHeap top=5); insert 15 (minHeap top=15); insert 1 ≤5 (maxHeap top=5, contains 5,1); insert 3 ≤5 (maxHeap top=5, contains 5,1,3), balance: len=3 >1+1, pop 5 -> push minHeap top=5 (15,5), maxHeap top=3 (1,3).

2. Step 2: Calculate median from heaps

   Heaps sizes equal (2 each). Median = average of maxHeap top (3) and minHeap top (5) = (3+5)/2 = 4.

3. Final Answer:

   4 -> Option D

4. Quick Check:

   Median = (3+5)/2 = 4 [OK]

Hint: Median is average of two heap tops if sizes equal [OK]

Common Mistakes:

• Taking median as maxHeap top only
• Not balancing heaps after insertion
• Calculating median as sum instead of average

4. In the following Go code snippet for median calculation, what is the bug?

func addNum(num int) {
  if maxHeap.Len() == 0 || num <= (*maxHeap)[0] {
    heap.Push(maxHeap, num)
  } else {
    heap.Push(minHeap, num)
  }
  if maxHeap.Len() > minHeap.Len()+1 {
    heap.Push(minHeap, heap.Pop(maxHeap))
  } else if minHeap.Len() > maxHeap.Len() {
    heap.Push(maxHeap, heap.Pop(minHeap))
  }
}

medium

A. Heaps are not balanced after insertion

B. Comparison operator should be < instead of <=

C. heap.Pop returns interface{}, needs type assertion before Push

D. maxHeap and minHeap are swapped in conditions

Solution

1. Step 1: Analyze balancing logic

   The first condition correctly balances when maxHeap is too large (> minHeap.Len()+1). But the second condition checks minHeap.Len() > maxHeap.Len(), which triggers rebalance correctly when minHeap is larger.

2. Step 2: Why heaps are not balanced

   Original code used > maxHeap.Len() + 1 which fails to rebalance when minHeap is exactly one larger. Changing to > maxHeap.Len() fixes this.

3. Final Answer:

   Heaps are not balanced after insertion -> Option A

4. Quick Check:

   minHeap balance: > max.Len() [OK]

Hint: minHeap balance when len(min) > len(max) [OK]

Common Mistakes:

• Using +1 symmetrically for both balance conditions
• Wrong size difference in minHeap check
• Not balancing after every insertion

5. You want to find the median of a large data stream where numbers can be negative, zero, or positive. Which approach using two heaps ensures correct median calculation even with duplicates and negative values?

hard

A. Use a single max-heap and pop elements until median is found

B. Use a max-heap for lower half and min-heap for upper half, balance sizes, and calculate median from tops

C. Sort all numbers in an array and pick middle element(s) each time

D. Use two min-heaps and merge their outputs to find median

Solution

1. Step 1: Choose correct heap types for halves

   Max-heap stores lower half (including negatives), min-heap stores upper half, handling all values correctly.

2. Step 2: Balance heaps and compute median

   Keep heaps balanced in size; median is top of max-heap if odd count, or average of tops if even count.

3. Final Answer:

   Use a max-heap for lower half and min-heap for upper half, balance sizes, and calculate median from tops -> Option B

4. Quick Check:

   Two heaps balanced handle negatives and duplicates [OK]

Hint: Balance max and min heaps for all number ranges [OK]

Common Mistakes:

• Using two min-heaps which can't split halves properly
• Sorting entire stream each time (inefficient)
• Using single heap only