Find Median from Data Stream

Hard

The problem

Build a MedianFinder. AddNum(num) adds a number to the collection, and FindMedian() returns the median of everything added so far: the middle value when sorted, or the average of the two middle values if the count is even.

Example 1
Input: AddNum(1) AddNum(2) FindMedian() AddNum(3) FindMedian()
Output: 1.5, 2.0
With [1, 2] the two middle values are 1 and 2, so the median is 1.5. With [1, 2, 3] the middle value is 2.

Limits

-100,000 ≤ num ≤ 100,000
FindMedian is only called when at least one number was added
Up to 50,000 calls in total

Write it in Go. Try for about 20 minutes on paper first, then open one hint at a time.

Try it here

Write Go. Common packages like fmt and sort are imported for you. Keep the function name and inputs the same.

Hints, one at a time

Nudge

The median sits between the smaller half and the larger half. Which structure gives you the edge of each half cheaply?

The idea

Max-heap for the lower half, min-heap for the upper half; keep sizes within 1 of each other. Median is the larger heap's top or the average of both tops.

Target: add O(log n), median O(1)

Go function shape

type MedianFinder struct{}
func Constructor() MedianFinder
func (m *MedianFinder) AddNum(num int)
func (m *MedianFinder) FindMedian() float64

Reference solution

Tested with go test. Try it yourself first, then compare.

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 any)        { *h = append(*h, x.(int)) }
func (h *maxHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}

// MedianFinder: a max-heap holds the smaller half, a min-heap the larger half.
// Invariant: len(lo) == len(hi) or len(lo) == len(hi)+1, and every lo value <= every hi value.
type MedianFinder struct {
	lo maxHeap // smaller half, root = its largest
	hi IntHeap // larger half, root = its smallest
}

func (m *MedianFinder) AddNum(x int) {
	heap.Push(&m.lo, x)
	heap.Push(&m.hi, heap.Pop(&m.lo).(int)) // pass the largest of lo across so halves stay ordered
	if m.hi.Len() > m.lo.Len() {
		heap.Push(&m.lo, heap.Pop(&m.hi).(int)) // rebalance sizes
	}
}

func (m *MedianFinder) FindMedian() float64 {
	if m.lo.Len() == 0 {
		return 0
	}
	if m.lo.Len() > m.hi.Len() {
		return float64(m.lo[0])
	}
	return float64(m.lo[0]+m.hi[0]) / 2
}

Read the lesson: Heap / Top-K