Find Median from Data Stream

Q: What is the time complexity of Find Median from Data Stream?

O(log n) per insertion, O(1) per median retrieval. Each insertion involves pushing into a heap and possibly rebalancing, each operation costing O(log n). Median retrieval is O(1) by peeking at the heap tops.

Q: What pattern does Find Median from Data Stream use?

The Pattern: Heap / Priority Queue (Dual Heap Median Maintenance) The Heuristic: Maintain two heaps to represent the lower and upper halves of the data stream, balancing their sizes to ensure the median can be retrieved in O(1) time, while insertions remain O(log n).

Hard Heaps

Problem

Design a data structure that supports adding numbers from a data stream and finding the median of all elements added so far.

At most 5 * 10⁴ calls will be made to addNum and findMedian.
−10⁵ ≤ num ≤ 10⁵

Example

Input: addNum(1), addNum(2), findMedian(), addNum(3), findMedian()

Output: 1.5, 2

After adding 1 and 2, the data structure contains [1, 2]. The median is (1 + 2) / 2 = 1.5. After adding 3, the data structure contains [1, 2, 3]. The median is 2, the middle element.

Approach

Straightforward Solution

A naive approach is to store all numbers in a list and sort it after each insertion. This leads to O(n log n) time per insertion, which is inefficient for large data streams.

Core Observation

The median divides a sorted list into two halves: the lower half and the upper half. Maintaining these halves separately allows efficient median retrieval. The lower half contains the smaller numbers, and the upper half contains the larger numbers. The median is either the maximum of the lower half, the minimum of the upper half, or the average of both, depending on the total number of elements.

Path to Optimal

The key recognition signals are 'median' and 'data stream', which indicate the need for a data structure that supports efficient insertion and median retrieval. Using two heaps (priority queues) — a max-heap for the lower half and a min-heap for the upper half — allows maintaining the halves in O(log n) time per insertion. Balancing the heaps ensures their sizes differ by at most one, enabling constant-time median calculation.

Optimal Approach

Use two heaps: a max-heap to store the smaller half of the numbers and a min-heap to store the larger half. When adding a number, push it into the max-heap (inverted sign to simulate max-heap). If the largest element in the max-heap is greater than the smallest in the min-heap, swap elements to maintain order. Then rebalance the heaps so their sizes differ by at most one. The median is the top of the larger heap if sizes differ, or the average of the tops if equal.

Time

O(log n) per insertion, O(1) per median retrieval

Each insertion involves pushing into a heap and possibly rebalancing, each operation costing O(log n). Median retrieval is O(1) by peeking at the heap tops.

Space

O(n)

All inserted elements are stored in the two heaps, so space grows linearly with the number of elements.

Pattern Spotlight

Heap / Priority Queue (Dual Heap Median Maintenance)

Maintain two heaps to represent the lower and upper halves of the data stream, balancing their sizes to ensure the median can be retrieved in O(1) time, while insertions remain O(log n).

Solution

1import heapq
2
3class MedianFinder:
  def __init__(self):
      self.small = []
      self.large = []
7
  def addNum(self, num: int) -> None:
      heapq.heappush(self.small, -1 * num)
      
      if (self.small and self.large and 
          (-1 * self.small[0]) > self.large[0]):
          val = -1 * heapq.heappop(self.small)
          heapq.heappush(self.large, val)
      
      if len(self.small) > len(self.large) + 1:
          val = -1 * heapq.heappop(self.small)
          heapq.heappush(self.large, val)
      if len(self.large) > len(self.small) + 1:
          val = heapq.heappop(self.large)
          heapq.heappush(self.small, -1 * val)
22
  def findMedian(self) -> float:
      if len(self.small) > len(self.large):
          return -1 * self.small[0]
      if len(self.large) > len(self.small):
          return self.large[0]
      
      return (-1 * self.small[0] + self.large[0]) / 2.0

1import java.util.PriorityQueue;
2import java.util.Collections;
3
4public class MedianFinder {
  private PriorityQueue<Integer> small;
  private PriorityQueue<Integer> large;
7
  public MedianFinder() {
      small = new PriorityQueue<>(Collections.reverseOrder());
      large = new PriorityQueue<>();
  }
12
  public void addNum(int num) {
      small.offer(num);
15
      if (!small.isEmpty() && !large.isEmpty() && small.peek() > large.peek()) {
          int val = small.poll();
          large.offer(val);
      }
20
      if (small.size() > large.size() + 1) {
          int val = small.poll();
          large.offer(val);
      }
      if (large.size() > small.size() + 1) {
          int val = large.poll();
          small.offer(val);
      }
  }
30
  public double findMedian() {
      if (small.size() > large.size()) {
          return small.peek();
      }
      if (large.size() > small.size()) {
          return large.peek();
      }
38
      return (small.peek() + large.peek()) / 2.0;
  }
41}

1#include <queue>
2
3using namespace std;
4
5class MedianFinder {
6private:
7    priority_queue<int> small; // max-heap
8    priority_queue<int, vector<int>, greater<int>> large; // min-heap
9
10public:
11    MedianFinder() {
12    }
13
14    void addNum(int num) {
15        small.push(num);
16
17        if (!small.empty() && !large.empty() && small.top() > large.top()) {
18            int val = small.top();
19            small.pop();
20            large.push(val);
21        }
22
23        if ((int)small.size() > (int)large.size() + 1) {
24            int val = small.top();
25            small.pop();
26            large.push(val);
27        }
28        if ((int)large.size() > (int)small.size() + 1) {
29            int val = large.top();
30            large.pop();
31            small.push(val);
32        }
33    }
34
35    double findMedian() {
36        if (small.size() > large.size()) {
37            return small.top();
38        }
39        if (large.size() > small.size()) {
40            return large.top();
41        }
42
43        return (small.top() + large.top()) / 2.0;
44    }
45};

1class MedianFinder {
  constructor() {
      this.small = new MaxHeap();
      this.large = new MinHeap();
  }
6
  addNum(num) {
      this.small.push(num);
9
      if (this.small.size() > 0 && this.large.size() > 0 && this.small.peek() > this.large.peek()) {
          const val = this.small.pop();
          this.large.push(val);
      }
14
      if (this.small.size() > this.large.size() + 1) {
          const val = this.small.pop();
          this.large.push(val);
      }
19
      if (this.large.size() > this.small.size() + 1) {
          const val = this.large.pop();
          this.small.push(val);
      }
  }
25
  findMedian() {
      if (this.small.size() > this.large.size()) {
          return this.small.peek();
      }
      if (this.large.size() > this.small.size()) {
          return this.large.peek();
      }
33
      return (this.small.peek() + this.large.peek()) / 2.0;
  }
36}
37
38// Helper MinHeap and MaxHeap implementations using arrays
39class MinHeap {
  constructor() {
      this.heap = [];
  }
43
  size() {
      return this.heap.length;
  }
47
  peek() {
      return this.heap[0];
  }
51
  push(val) {
      this.heap.push(val);
      this._bubbleUp(this.heap.length - 1);
  }
56
  pop() {
      if (this.heap.length === 1) return this.heap.pop();
      const top = this.heap[0];
      this.heap[0] = this.heap.pop();
      this._bubbleDown(0);
      return top;
  }
64
  _bubbleUp(index) {
      while (index > 0) {
          const parent = Math.floor((index - 1) / 2);
          if (this.heap[parent] <= this.heap[index]) break;
          [this.heap[parent], this.heap[index]] = [this.heap[index], this.heap[parent]];
          index = parent;
      }
  }
73
  _bubbleDown(index) {
      const length = this.heap.length;
      while (true) {
          let left = 2 * index + 1;
          let right = 2 * index + 2;
          let smallest = index;
80
          if (left < length && this.heap[left] < this.heap[smallest]) smallest = left;
          if (right < length && this.heap[right] < this.heap[smallest]) smallest = right;
83
          if (smallest === index) break;
85
          [this.heap[smallest], this.heap[index]] = [this.heap[index], this.heap[smallest]];
          index = smallest;
      }
  }
90}
91
92class MaxHeap {
  constructor() {
      this.heap = [];
  }
96
  size() {
      return this.heap.length;
  }
100
  peek() {
      return this.heap[0];
  }
104
  push(val) {
      this.heap.push(val);
      this._bubbleUp(this.heap.length - 1);
  }
109
  pop() {
      if (this.heap.length === 1) return this.heap.pop();
      const top = this.heap[0];
      this.heap[0] = this.heap.pop();
      this._bubbleDown(0);
      return top;
  }
117
  _bubbleUp(index) {
      while (index > 0) {
          const parent = Math.floor((index - 1) / 2);
          if (this.heap[parent] >= this.heap[index]) break;
          [this.heap[parent], this.heap[index]] = [this.heap[index], this.heap[parent]];
          index = parent;
      }
  }
126
  _bubbleDown(index) {
      const length = this.heap.length;
      while (true) {
          let left = 2 * index + 1;
          let right = 2 * index + 2;
          let largest = index;
133
          if (left < length && this.heap[left] > this.heap[largest]) largest = left;
          if (right < length && this.heap[right] > this.heap[largest]) largest = right;
136
          if (largest === index) break;
138
          [this.heap[largest], this.heap[index]] = [this.heap[index], this.heap[largest]];
          index = largest;
      }
  }
143}

Step-by-Step Solution

Initialize Two Heaps for Median Maintenance

4def __init__(self):
5    self.small = []
6    self.large = []

Objective

To create the lower-half max-heap and upper-half min-heap used to maintain the stream median.

Key Insight

The median is determined by the boundary between the lower and upper halves. Keeping the lower half in a max-heap and the upper half in a min-heap gives constant-time access to the two middle candidates.

Interview Quick-Check

Core Logic

The lower half needs max-heap behavior; Python simulates it with negative values.

Core Logic

The upper half uses a normal min-heap.

Insert Number and Maintain Heap Order

8def addNum(self, num: int) -> None:
  heapq.heappush(self.small, -1 * num)
  if (self.small and self.large and 
      (-1 * self.small[0]) > self.large[0]):
      val = -1 * heapq.heappop(self.small)
      heapq.heappush(self.large, val)

Objective

To insert the new number and restore the invariant that every lower-half value is less than or equal to every upper-half value.

Key Insight

The new number is first added to the lower half. If the largest lower-half value is greater than the smallest upper-half value, that boundary value is moved into the upper half.

Interview Quick-Check

Core Logic

Insert into the lower heap first, then repair ordering if the heap tops cross.

State & Boundaries

Compare only after confirming both heaps contain values.

Balance Heap Sizes to Maintain Median Property

16if len(self.small) > len(self.large) + 1:
17    val = -1 * heapq.heappop(self.small)
18    heapq.heappush(self.large, val)
19if len(self.large) > len(self.small) + 1:
20    val = heapq.heappop(self.large)
21    heapq.heappush(self.small, -1 * val)

Objective

To ensure the two heaps differ in size by at most one element.

Key Insight

The median is either the top of the larger heap or the average of both heap tops. That only works if the heap sizes remain balanced.

Interview Quick-Check

Core Logic

Move one top element from the larger heap to the smaller heap when size difference exceeds one.

Retrieve Median Based on Heap Sizes

23def findMedian(self) -> float:
  if len(self.small) > len(self.large):
      return -1 * self.small[0]
  if len(self.large) > len(self.small):
      return self.large[0]
  return (-1 * self.small[0] + self.large[0]) / 2.0

Objective

To return the correct median from heap tops in constant time.

Key Insight

If one heap is larger, its top is the median. If both heaps are equal size, the median is the average of the lower-half maximum and upper-half minimum.

Interview Quick-Check

Common Pitfalls & Bugs

Remember to negate the lower heap top in Python before using it as a real value.

Line Analysis

This solution has 4 Critical lines interviewers watch for.

Line 9 Critical

heapq.heappush(self.small, -1 * num)

Push the negative of the new number into the max-heap to simulate max-heap behavior.

Python's heapq is a min-heap by default; negating the number allows it to behave as a max-heap, ensuring the largest number in the lower half is accessible at the top.

Line 12 Critical

(-1 * self.small[0]) > self.large[0]):

Compare the lower-half maximum against the upper-half minimum.

If the largest lower-half value exceeds the smallest upper-half value, the two heaps are ordered incorrectly and must be repaired.

Line 25 Critical

return -1 * self.small[0]

Return the median as the negated top element of the max-heap.

Negating the stored negative value retrieves the actual median when the lower half is larger.

Line 29 Critical

return (-1 * self.small[0] + self.large[0]) / 2.0

Return the average of the tops of both heaps when they have equal size.

When both halves are equal in size, the median is the average of the largest element in the lower half and the smallest in the upper half, computed efficiently by accessing the heap tops.

Line 5 Core

self.small = []

Initialize an empty max-heap to store the smaller half of the numbers.

Using a max-heap for the lower half allows quick access to the largest element in the smaller half, which is essential for median calculation and maintaining order between halves.

Line 6 Core

self.large = []

Initialize an empty min-heap to store the larger half of the numbers.

The min-heap stores the larger half, enabling quick access to the smallest element in the upper half, which is critical for balancing and median retrieval.

Line 11 Core

if (self.small and self.large and

Check whether both heaps exist before comparing their boundary values.

This condition starts the order-invariant check only when both heaps have values to compare.

Line 13 Core

val = -1 * heapq.heappop(self.small)

Remove the largest element from the max-heap.

Popping this element prepares it for insertion into the min-heap, restoring the order invariant between halves.

Line 14 Core

heapq.heappush(self.large, val)

Push the extracted value into the min-heap.

Inserting into the min-heap places the element in the upper half, maintaining the correct partitioning of the data stream.

Line 16 Core

if len(self.small) > len(self.large) + 1:

Check if the max-heap is larger than the min-heap by more than one element.

This condition identifies imbalance in heap sizes that would invalidate median calculation, triggering rebalancing.

Line 17 Core

val = -1 * heapq.heappop(self.small)

Extract the largest element from the max-heap for rebalancing.

Removing the top element from the larger heap prepares it for transfer to the smaller heap to restore balance.

Line 18 Core

heapq.heappush(self.large, val)

Insert the extracted element into the min-heap.

Moving the element to the smaller heap restores the size invariant, ensuring median correctness.

Line 19 Core

if len(self.large) > len(self.small) + 1:

Check if the min-heap is larger than the max-heap by more than one element.

This condition detects imbalance in the opposite direction, requiring rebalancing to maintain the median property.

Line 20 Core

val = heapq.heappop(self.large)

Extract the smallest element from the min-heap for rebalancing.

Removing the top element from the larger heap prepares it for transfer to the smaller heap to restore balance.

Line 21 Core

heapq.heappush(self.small, -1 * val)

Insert the negated extracted element into the max-heap.

Negating the value before insertion maintains the max-heap property of the lower half, completing the rebalance.

Line 27 Core

return self.large[0]

Return the median as the top element of the min-heap.

Directly returning the min-heap's top element yields the median when the upper half is larger.

Line 24 Core

if len(self.small) > len(self.large):

Check if the max-heap contains more elements than the min-heap.

If true, the median is the largest element in the lower half, accessible at the top of the max-heap.

Line 26 Core

if len(self.large) > len(self.small):

Check if the min-heap contains more elements than the max-heap.

If true, the median is the smallest element in the upper half, accessible at the top of the min-heap.

Line 4 Standard

def __init__(self):

Define the MedianFinder constructor.

The constructor initializes the two heap structures used throughout the data stream.

Line 8 Standard

def addNum(self, num: int) -> None:

Define the addNum method.

This method inserts a new stream value while preserving median-maintenance invariants.

Line 23 Standard

def findMedian(self) -> float:

Define the findMedian method.

This method reads the maintained heap state and returns the current stream median.

Test Your Understanding

Why are two heaps necessary instead of a single heap or a sorted list?

Two heaps allow splitting the data into lower and upper halves, enabling efficient insertion and median retrieval. A single heap cannot provide both the maximum of the lower half and minimum of the upper half efficiently, and a sorted list requires O(n) insertion time.

Don't just read it. Drill it.

Reconstruct Find Median from Data Stream from memory until it sticks. AlgoDrill blanks out key lines and makes you fill them back in, step by step.

Unlock the Find Median from Data Stream drill

or drill a free problem

1	`import heapq`
2
3	`class MedianFinder:`
4	`def __init__(self):`
5	`self.small = []`
6	`self.large = []`
7
8	`def addNum(self, num: int) -> None:`
9	`heapq.heappush(self.small, -1 * num)`
10
11	`if (self.small and self.large and`
12	`(-1 * self.small[0]) > self.large[0]):`
13	`val = -1 * heapq.heappop(self.small)`
14	`heapq.heappush(self.large, val)`
15
16	`if len(self.small) > len(self.large) + 1:`
17	`val = -1 * heapq.heappop(self.small)`
18	`heapq.heappush(self.large, val)`
19	`if len(self.large) > len(self.small) + 1:`
20	`val = heapq.heappop(self.large)`
21	`heapq.heappush(self.small, -1 * val)`
22
23	`def findMedian(self) -> float:`
24	`if len(self.small) > len(self.large):`
25	`return -1 * self.small[0]`
26	`if len(self.large) > len(self.small):`
27	`return self.large[0]`
28
29	`return (-1 * self.small[0] + self.large[0]) / 2.0`

1	`import java.util.PriorityQueue;`
2	`import java.util.Collections;`
3
4	`public class MedianFinder {`
5	`private PriorityQueue<Integer> small;`
6	`private PriorityQueue<Integer> large;`
7
8	`public MedianFinder() {`
9	`small = new PriorityQueue<>(Collections.reverseOrder());`
10	`large = new PriorityQueue<>();`
11	`}`
12
13	`public void addNum(int num) {`
14	`small.offer(num);`
15
16	`if (!small.isEmpty() && !large.isEmpty() && small.peek() > large.peek()) {`
17	`int val = small.poll();`
18	`large.offer(val);`
19	`}`
20
21	`if (small.size() > large.size() + 1) {`
22	`int val = small.poll();`
23	`large.offer(val);`
24	`}`
25	`if (large.size() > small.size() + 1) {`
26	`int val = large.poll();`
27	`small.offer(val);`
28	`}`
29	`}`
30
31	`public double findMedian() {`
32	`if (small.size() > large.size()) {`
33	`return small.peek();`
34	`}`
35	`if (large.size() > small.size()) {`
36	`return large.peek();`
37	`}`
38
39	`return (small.peek() + large.peek()) / 2.0;`
40	`}`
41	`}`

1	`#include <queue>`
2
3	`using namespace std;`
4
5	`class MedianFinder {`
6	`private:`
7	`priority_queue<int> small; // max-heap`
8	`priority_queue<int, vector<int>, greater<int>> large; // min-heap`
9
10	`public:`
11	`MedianFinder() {`
12	`}`
13
14	`void addNum(int num) {`
15	`small.push(num);`
16
17	`if (!small.empty() && !large.empty() && small.top() > large.top()) {`
18	`int val = small.top();`
19	`small.pop();`
20	`large.push(val);`
21	`}`
22
23	`if ((int)small.size() > (int)large.size() + 1) {`
24	`int val = small.top();`
25	`small.pop();`
26	`large.push(val);`
27	`}`
28	`if ((int)large.size() > (int)small.size() + 1) {`
29	`int val = large.top();`
30	`large.pop();`
31	`small.push(val);`
32	`}`
33	`}`
34
35	`double findMedian() {`
36	`if (small.size() > large.size()) {`
37	`return small.top();`
38	`}`
39	`if (large.size() > small.size()) {`
40	`return large.top();`
41	`}`
42
43	`return (small.top() + large.top()) / 2.0;`
44	`}`
45	`};`

1	`class MedianFinder {`
2	`constructor() {`
3	`this.small = new MaxHeap();`
4	`this.large = new MinHeap();`
5	`}`
6
7	`addNum(num) {`
8	`this.small.push(num);`
9
10	`if (this.small.size() > 0 && this.large.size() > 0 && this.small.peek() > this.large.peek()) {`
11	`const val = this.small.pop();`
12	`this.large.push(val);`
13	`}`
14
15	`if (this.small.size() > this.large.size() + 1) {`
16	`const val = this.small.pop();`
17	`this.large.push(val);`
18	`}`
19
20	`if (this.large.size() > this.small.size() + 1) {`
21	`const val = this.large.pop();`
22	`this.small.push(val);`
23	`}`
24	`}`
25
26	`findMedian() {`
27	`if (this.small.size() > this.large.size()) {`
28	`return this.small.peek();`
29	`}`
30	`if (this.large.size() > this.small.size()) {`
31	`return this.large.peek();`
32	`}`
33
34	`return (this.small.peek() + this.large.peek()) / 2.0;`
35	`}`
36	`}`
37
38	`// Helper MinHeap and MaxHeap implementations using arrays`
39	`class MinHeap {`
40	`constructor() {`
41	`this.heap = [];`
42	`}`
43
44	`size() {`
45	`return this.heap.length;`
46	`}`
47
48	`peek() {`
49	`return this.heap[0];`
50	`}`
51
52	`push(val) {`
53	`this.heap.push(val);`
54	`this._bubbleUp(this.heap.length - 1);`
55	`}`
56
57	`pop() {`
58	`if (this.heap.length === 1) return this.heap.pop();`
59	`const top = this.heap[0];`
60	`this.heap[0] = this.heap.pop();`
61	`this._bubbleDown(0);`
62	`return top;`
63	`}`
64
65	`_bubbleUp(index) {`
66	`while (index > 0) {`
67	`const parent = Math.floor((index - 1) / 2);`
68	`if (this.heap[parent] <= this.heap[index]) break;`
69	`[this.heap[parent], this.heap[index]] = [this.heap[index], this.heap[parent]];`
70	`index = parent;`
71	`}`
72	`}`
73
74	`_bubbleDown(index) {`
75	`const length = this.heap.length;`
76	`while (true) {`
77	`let left = 2 * index + 1;`
78	`let right = 2 * index + 2;`
79	`let smallest = index;`
80
81	`if (left < length && this.heap[left] < this.heap[smallest]) smallest = left;`
82	`if (right < length && this.heap[right] < this.heap[smallest]) smallest = right;`
83
84	`if (smallest === index) break;`
85
86	`[this.heap[smallest], this.heap[index]] = [this.heap[index], this.heap[smallest]];`
87	`index = smallest;`
88	`}`
89	`}`
90	`}`
91
92	`class MaxHeap {`
93	`constructor() {`
94	`this.heap = [];`
95	`}`
96
97	`size() {`
98	`return this.heap.length;`
99	`}`
100
101	`peek() {`
102	`return this.heap[0];`
103	`}`
104
105	`push(val) {`
106	`this.heap.push(val);`
107	`this._bubbleUp(this.heap.length - 1);`
108	`}`
109
110	`pop() {`
111	`if (this.heap.length === 1) return this.heap.pop();`
112	`const top = this.heap[0];`
113	`this.heap[0] = this.heap.pop();`
114	`this._bubbleDown(0);`
115	`return top;`
116	`}`
117
118	`_bubbleUp(index) {`
119	`while (index > 0) {`
120	`const parent = Math.floor((index - 1) / 2);`
121	`if (this.heap[parent] >= this.heap[index]) break;`
122	`[this.heap[parent], this.heap[index]] = [this.heap[index], this.heap[parent]];`
123	`index = parent;`
124	`}`
125	`}`
126
127	`_bubbleDown(index) {`
128	`const length = this.heap.length;`
129	`while (true) {`
130	`let left = 2 * index + 1;`
131	`let right = 2 * index + 2;`
132	`let largest = index;`
133
134	`if (left < length && this.heap[left] > this.heap[largest]) largest = left;`
135	`if (right < length && this.heap[right] > this.heap[largest]) largest = right;`
136
137	`if (largest === index) break;`
138
139	`[this.heap[largest], this.heap[index]] = [this.heap[index], this.heap[largest]];`
140	`index = largest;`
141	`}`
142	`}`
143	`}`

8	`def addNum(self, num: int) -> None:`
9	`heapq.heappush(self.small, -1 * num)`
11	`if (self.small and self.large and`
12	`(-1 * self.small[0]) > self.large[0]):`
13	`val = -1 * heapq.heappop(self.small)`
14	`heapq.heappush(self.large, val)`

Problem

Example

Approach

Solution

Step-by-Step Solution

Initialize Two Heaps for Median Maintenance

Insert Number and Maintain Heap Order

Balance Heap Sizes to Maintain Median Property

Retrieve Median Based on Heap Sizes

Line Analysis

Test Your Understanding

Related Problems

Don't just read it. Drill it.