Introduction to Find Median from Data Stream
The "Find Median from Data Stream" problem is a classic algorithmic challenge frequently encountered in coding interviews and real-world data processing applications. The problem asks you to design a data structure that efficiently supports two operations: adding numbers from a continuous stream of data, and finding the median of all numbers added so far at any point in time.
The median of a sorted dataset is the middle value if the count is odd, or the average of the two middle values if the count is even. While computing the median of a static array is straightforward, doing so efficiently on a continuously growing stream requires careful data structure selection.
Why This Problem Matters
Streaming median computation is essential in many domains. In financial systems, traders monitor median stock prices in real time. In monitoring and observability tools, median latency is a more robust metric than average latency because it is resistant to outliers. In machine learning pipelines, streaming medians help with feature normalization and anomaly detection on incoming batches of data.
A naive approach — sorting the entire dataset each time a new number arrives — yields O(n log n) time per query, which becomes prohibitively expensive as the stream grows. The heap-based solution we will explore reduces this to O(log n) per insertion and O(1) per median query, making it suitable for high-throughput systems.
Understanding the Approach
The key insight is to maintain two heaps that partition the incoming numbers:
- Max-heap for the lower half of the numbers. The largest of the smaller half sits at the top.
- Min-heap for the upper half of the numbers. The smallest of the larger half sits at the top.
By keeping the two heaps balanced (their sizes differ by at most one), the median can always be derived from the top elements of these heaps. If the heaps are equal in size, the median is the average of both tops. Otherwise, the median is the top of the heap that contains the extra element.
Why Go?
Go's standard library provides the container/heap package, which implements a generic heap interface. Combined with Go's strong typing, concurrency primitives, and excellent performance characteristics, it is an ideal language for building streaming data structures that may be used in production services.
Implementing the Heaps in Go
Go's container/heap package requires us to implement the heap.Interface, which consists of sort.Interface plus Push and Pop methods. Let's start by defining a max-heap of integers.
package medianstream
import "container/heap"
// MaxHeap stores integers in descending order.
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{} {
old := *h
n := len(old)
x := old[n-1]
*h = old[:n-1]
return x
}
Next, we define a min-heap. The only difference is the Less function, which compares in the opposite direction.
// MinHeap stores integers in ascending order.
type MinHeap []int
func (h MinHeap) Len() int { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x interface{}) {
*h = append(*h, x.(int))
}
func (h *MinHeap) Pop() interface{} {
old := *h
n := len(old)
x := old[n-1]
*h = old[:n-1]
return x
}
Building the MedianFinder
With both heaps defined, we can now build the MedianFinder struct. It holds references to both heaps and exposes AddNum and FindMedian methods.
// MedianFinder maintains a stream of numbers and supports
// efficient median queries.
type MedianFinder struct {
lower *MaxHeap // stores the smaller half
upper *MinHeap // stores the larger half
}
// NewMedianFinder constructs and initializes a MedianFinder.
func NewMedianFinder() *MedianFinder {
lower := &MaxHeap{}
upper := &MinHeap{}
heap.Init(lower)
heap.Init(upper)
return &MedianFinder{lower: lower, upper: upper}
}
Adding Numbers
The AddNum method must maintain two invariants: every element in lower must be less than or equal to every element in upper, and the sizes of the two heaps must differ by at most one. We achieve this with a careful insertion strategy.
// AddNum inserts a number into the stream.
func (mf *MedianFinder) AddNum(num int) {
// Step 1: Decide which heap to insert into.
if mf.lower.Len() == 0 || num <= (*mf.lower)[0] {
heap.Push(mf.lower, num)
} else {
heap.Push(mf.upper, num)
}
// Step 2: Rebalance the heaps so their sizes differ by at most 1.
if mf.lower.Len() > mf.upper.Len()+1 {
heap.Push(mf.upper, heap.Pop(mf.lower))
} else if mf.upper.Len() > mf.lower.Len() {
heap.Push(mf.lower, heap.Pop(mf.upper))
}
}
The logic is straightforward. We first place the new number into the appropriate heap based on its value relative to the current maximum of the lower half. Then, if one heap becomes too large, we move its top element to the other heap. This guarantees the balance invariant is preserved after every insertion.
Finding the Median
Once the heaps are balanced, finding the median is an O(1) operation.
// FindMedian returns the median of all numbers added so far.
func (mf *MedianFinder) FindMedian() float64 {
if mf.lower.Len() > mf.upper.Len() {
return float64((*mf.lower)[0])
}
if mf.upper.Len() > mf.lower.Len() {
return float64((*mf.upper)[0])
}
// Equal sizes: average the two tops.
return float64((*mf.lower)[0]+(*mf.upper)[0]) / 2.0
}
Putting It All Together
Here is the complete, runnable program including a main function that demonstrates the median finder in action.
package main
import (
"container/heap"
"fmt"
)
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{} {
old := *h
n := len(old)
x := old[n-1]
*h = old[:n-1]
return x
}
type MinHeap []int
func (h MinHeap) Len() int { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x interface{}) {
*h = append(*h, x.(int))
}
func (h *MinHeap) Pop() interface{} {
old := *h
n := len(old)
x := old[n-1]
*h = old[:n-1]
return x
}
type MedianFinder struct {
lower *MaxHeap
upper *MinHeap
}
func NewMedianFinder() *MedianFinder {
lower := &MaxHeap{}
upper := &MinHeap{}
heap.Init(lower)
heap.Init(upper)
return &MedianFinder{lower: lower, upper: upper}
}
func (mf *MedianFinder) AddNum(num int) {
if mf.lower.Len() == 0 || num <= (*mf.lower)[0] {
heap.Push(mf.lower, num)
} else {
heap.Push(mf.upper, num)
}
if mf.lower.Len() > mf.upper.Len()+1 {
heap.Push(mf.upper, heap.Pop(mf.lower))
} else if mf.upper.Len() > mf.lower.Len() {
heap.Push(mf.lower, heap.Pop(mf.upper))
}
}
func (mf *MedianFinder) FindMedian() float64 {
if mf.lower.Len() > mf.upper.Len() {
return float64((*mf.lower)[0])
}
if mf.upper.Len() > mf.lower.Len() {
return float64((*mf.upper)[0])
}
return float64((*mf.lower)[0]+(*mf.upper)[0]) / 2.0
}
func main() {
mf := NewMedianFinder()
values := []int{5, 15, 1, 3, 8, 7, 9, 2, 6}
for _, v := range values {
mf.AddNum(v)
fmt.Printf("Added %d, median so far: %.2f\n", v, mf.FindMedian())
}
}
Running this program produces output showing the median evolving as each number is added. After inserting 5, the median is 5. After inserting 15, the median becomes 10 (the average of 5 and 15). After inserting 1, the median is 5, and so on. This demonstrates that the structure correctly maintains the median across an arbitrary sequence of insertions.
Complexity Analysis
Understanding the performance characteristics of this solution is critical for evaluating its suitability in production systems.
- AddNum: O(log n) time, because each insertion involves at most one heap push and one heap pop, both of which are logarithmic in the size of the heap.
- FindMedian: O(1) time, since we only read the top elements of the heaps.
- Space: O(n), where n is the number of elements added, since every element is stored in exactly one of the two heaps.
Compare this to the naive approach of sorting the entire array on each query, which would cost O(n log n) per median query. For a stream of one million numbers, the heap-based approach is dramatically faster in practice.
Best Practices
Use Generics for Reusability
If you are using Go 1.18 or later, consider rewriting the heaps using type parameters so the same code can compute streaming medians for float64, int64, or any ordered type. This avoids code duplication and makes the utility more broadly applicable across your codebase.
Guard Against Empty State
In production code, calling FindMedian before any numbers have been added is an edge case that should be handled explicitly. Returning an error or a sentinel value such as math.NaN() is safer than returning zero, which could be mistaken for a legitimate median.
import "math"
func (mf *MedianFinder) FindMedian() (float64, error) {
if mf.lower.Len() == 0 && mf.upper.Len() == 0 {
return math.NaN(), fmt.Errorf("no numbers added yet")
}
// ... rest of the logic
}
Consider Concurrency
If the median finder is shared across goroutines, wrap the AddNum and FindMedian methods with a sync.Mutex. Alternatively, use a channel-based design where a single goroutine owns the structure and processes commands from other goroutines, which avoids lock contention under heavy load.
Validate Inputs
When integrating with external data sources, validate incoming numbers before adding them. NaN or infinite float values can corrupt the heap invariants and produce incorrect medians. Rejecting invalid inputs early prevents subtle bugs from propagating downstream.
Benchmark Before Optimizing
The heap-based solution is already optimal in asymptotic terms, but micro-optimizations such as pre-allocating slice capacity or avoiding the interface{} conversions in Push and Pop can yield measurable gains in hot paths. Always benchmark with realistic data before committing to such changes.
Conclusion
The two-heap technique for finding the median of a data stream is an elegant example of how choosing the right data structure transforms an expensive problem into an efficient one. By partitioning incoming numbers into a max-heap for the lower half and a min-heap for the upper half, we achieve logarithmic insertions and constant-time median queries, making the solution practical for high-volume streaming applications. Go's container/heap package provides a clean foundation for this design, and with attention to edge cases, concurrency, and input validation, the implementation can serve as a reliable component in production systems ranging from financial analytics to observability platforms.