-
Notifications
You must be signed in to change notification settings - Fork 0
100 — Find Median from Data Stream
LeetCode 295 · Hard. Design a data structure that supports addNum(_ num: Int) (adds an integer from a data stream) and findMedian() -> Double (returns the median of all elements added so far). findMedian may be called many times, interleaved with calls to addNum.
Picture a single-file line of people, sorted shortest to tallest, and you're asked "who's in the middle?" every time someone new joins. Re-sorting the entire line from scratch every time someone joins works, but it's overkill — what if, instead, you kept two separate huddles: a "shorter half" huddle where the tallest person always stands at the front, and a "taller half" huddle where the shortest person always stands at the front? Keep the two huddles the same size (or the shorter-half huddle exactly one person bigger). The median is then always readable directly off the two people standing at the front — no need to know anything about anyone standing further back in either huddle.
"Continuously find the median as data streams in" — any "design a class with addNum/findMedian" phrasing, or more generally any problem needing repeated access to the middle of a growing, unsorted collection, is the classic two-heap signature: a max-heap for the lower half, a min-heap for the upper half.
Keep every number seen so far in an unsorted array. addNum is a cheap append; findMedian sorts a fresh copy every single call and reads off the middle position(s).
class MedianFinderBruteForce {
private var nums: [Int] = []
func addNum(_ num: Int) {
nums.append(num)
}
func findMedian() -> Double {
let sorted = nums.sorted()
let n = sorted.count
if n % 2 == 1 {
return Double(sorted[n / 2])
} else {
return (Double(sorted[n / 2 - 1]) + Double(sorted[n / 2])) / 2.0
}
}
}
// smoke test — mirrors LeetCode's canonical example
let bfMedian = MedianFinderBruteForce()
bfMedian.addNum(1)
bfMedian.addNum(2)
print(bfMedian.findMedian()) // 1.5
bfMedian.addNum(3)
print(bfMedian.findMedian()) // 2.0
// edge case: repeated/tied values
let bfMedian2 = MedianFinderBruteForce()
bfMedian2.addNum(5)
bfMedian2.addNum(5)
bfMedian2.addNum(5)
print(bfMedian2.findMedian()) // 5.0Big-O: addNum is O(1) amortized. findMedian is O(n log n) — it re-sorts the entire stream history every single time it's called, even though only the one or two middle elements are ever read.
Split the stream across two heaps: maxHeap holds the smaller half of the numbers (its root is the largest of the small half), and minHeap holds the larger half (its root is the smallest of the large half). After every insert, rebalance so the two heaps differ in size by at most one, with maxHeap allowed to hold exactly one extra element when the total count is odd. The median is then always sitting at one or both roots.
struct Heap<T> {
private var elements: [T] = []
private let areInIncreasingOrder: (T, T) -> Bool
init(sort: @escaping (T, T) -> Bool) {
self.areInIncreasingOrder = sort
}
var isEmpty: Bool { elements.isEmpty }
var count: Int { elements.count }
var peek: T? { elements.first }
mutating func insert(_ value: T) {
elements.append(value)
siftUp(from: elements.count - 1)
}
mutating func extract() -> T? {
guard !elements.isEmpty else { return nil }
elements.swapAt(0, elements.count - 1)
let top = elements.removeLast()
siftDown(from: 0)
return top
}
private mutating func siftUp(from index: Int) {
var child = index
var parent = (child - 1) / 2
while child > 0 && areInIncreasingOrder(elements[child], elements[parent]) {
elements.swapAt(child, parent)
child = parent
parent = (child - 1) / 2
}
}
private mutating func siftDown(from index: Int) {
var parent = index
while true {
let left = 2 * parent + 1
let right = 2 * parent + 2
var candidate = parent
if left < elements.count && areInIncreasingOrder(elements[left], elements[candidate]) {
candidate = left
}
if right < elements.count && areInIncreasingOrder(elements[right], elements[candidate]) {
candidate = right
}
if candidate == parent { return }
elements.swapAt(parent, candidate)
parent = candidate
}
}
}
class MedianFinder {
private var maxHeap = Heap<Int>(sort: >) // lower half; largest of the lower half sits at the root
private var minHeap = Heap<Int>(sort: <) // upper half; smallest of the upper half sits at the root
func addNum(_ num: Int) {
if maxHeap.isEmpty || num <= maxHeap.peek! {
maxHeap.insert(num)
} else {
minHeap.insert(num)
}
// rebalance so the two halves differ in size by at most 1,
// with maxHeap allowed to hold exactly one more than minHeap
if maxHeap.count > minHeap.count + 1 {
minHeap.insert(maxHeap.extract()!)
} else if minHeap.count > maxHeap.count {
maxHeap.insert(minHeap.extract()!)
}
}
func findMedian() -> Double {
if maxHeap.count > minHeap.count {
return Double(maxHeap.peek!)
}
return (Double(maxHeap.peek!) + Double(minHeap.peek!)) / 2.0
}
}
// smoke test — same cases as the brute force
let median = MedianFinder()
median.addNum(1)
median.addNum(2)
print(median.findMedian()) // 1.5
median.addNum(3)
print(median.findMedian()) // 2.0
let median2 = MedianFinder()
median2.addNum(5)
print(median2.findMedian()) // 5.0
median2.addNum(5)
print(median2.findMedian()) // 5.0
median2.addNum(5)
print(median2.findMedian()) // 5.0Big-O: addNum is O(log n) — one heap insert plus, at most, one rebalancing extract/insert pair, all O(log n). findMedian is O(1) — just reading one or two roots. O(n) total space across both heaps.
The brute force's findMedian throws away all the sorting work it did last time and starts over, even though inserting one new number can shift the median by at most one position. The two-heap approach keeps that "almost sorted around the middle" structure alive permanently: everything below the median lives in a max-heap (so the largest of the small half — the left boundary of the median — is always the root), everything above lives in a min-heap (so the smallest of the large half — the right boundary — is always the root), and the size-balancing invariant guarantees those two roots are always exactly the elements the median formula needs. Moving the boundary when a new number arrives costs O(log n); reading it costs O(1) — a full re-sort is never needed again.
- Heaps and Priority Queues — this problem is the canonical two-heap technique: one max-heap, one min-heap, working together.
- Top K and Heap Pattern — explicitly calls out "a running/streaming median" as one of this pattern's core recognition signals, pointing straight at the two-heap setup used here.
- Median of Two Sorted Arrays — a different median problem worth contrasting: that one has two already-sorted, static arrays and reaches for binary search, whereas this one has a single growing, unsorted stream and reaches for two balanced heaps — same target quantity (the median), entirely different technique because the inputs' shape is different.
Find Median from Data Stream is two huddles facing each other across the median line — a max-heap fronts the short half, a min-heap fronts the tall half, and the median is always sitting at the two people up front.
- Why does
addNumcompare the new number againstmaxHeap.peek!(notminHeap.peek!) to decide which heap it goes into first? - Walk through
addNumon the sequence[5, 5, 5]by hand. Why does the size-balancing step in the second call move a value frommaxHeaptominHeap, given that both heaps only ever contain the same repeated value? - Why is
maxHeapallowed to hold exactly one more element thanminHeap, but never the other way around? What wouldfindMedianneed to look like if the reverse imbalance were allowed instead? - If
findMedianis called before any numbers have been added,maxHeap.peek!force-unwrapsniland crashes. Why does the problem's own constraints make this safe to leave unhandled here?