Find the Median from a Data Stream Using Two Heaps

59.1k
0

Design a class for a growing stream of integers. The class must support two operations:

- addNum(value): Add value to the stream. - findMedian(): Return the median of all values added so far.

An odd-sized ordered stream has one middle value. An even-sized ordered stream has a median equal to the average of the two middle values. Every findMedian call occurs after at least one insertion.

Example 1

Input: [Solution(), addNum(1), addNum(2), findMedian(), addNum(3), findMedian()]
Output: [null, null, null, 1.5, null, 2.0]
Explanation: Values [1, 2] have two middle positions, so the first median equals (1 + 2) / 2 = 1.5. Values [1, 2, 3] have middle value 2, so the second median equals 2.0.

Example 2

Input: [Solution(), addNum(-5), addNum(-5), findMedian(), addNum(10), findMedian()]
Output: [null, null, null, -5.0, null, -5.0]
Explanation: Duplicate values occupy separate positions. Ordered values [-5, -5] produce median -5.0, and ordered values [-5, -5, 10] also produce median -5.0.

Brute Force Approach

The easiest starting point saves every arriving value. No ordering work is needed during insertion, so each new number simply joins the stored stream.

A median query creates a sorted copy. Sorting exposes the middle position directly, but repeated queries rebuild the same order again and again. Long streams therefore make the simple method expensive.

Algorithm

  • Begin with an empty list of stream values, allowing each insertion to preserve the complete history without extra ordering work.

  • Append every new value to the list because a later median query depends on every occurrence, including duplicates.

  • Copy the stored list during findMedian, preserving the insertion storage while a temporary ordered view is prepared.

  • Sort the copied values in ascending order because direct middle indices define the median in an ordered collection.

  • Return index n / 2 for an odd size because zero-based indexing places the single middle value at index n / 2.

  • Average indices n / 2 - 1 and n / 2 for an even size, using two half-values to avoid integer overflow during addition.

Dry Run

Median from Data Stream - Brute Approach

Median from Data Stream - Brute Approach

Solution

#include <bits/stdc++.h>
using namespace std;
class Solution {
private:
vector<int> numbers;
public:
// Create empty stream storage.
Solution() {
}
// Add one value to the stored stream.
void addNum(int value) {
// Preserve the newest stream value.
numbers.push_back(value);
}
// Return the median after sorting a copied stream.
double findMedian() {
// Copy values so insertion storage stays unchanged.
vector<int> ordered = numbers;
// Sort the copy to expose the middle position.
sort(ordered.begin(), ordered.end());
int size = ordered.size();
int middle = size / 2;
// Odd size has one value at the middle index.
if (size % 2 == 1) {
return ordered[middle];
}
// Half-value addition avoids integer overflow.
return ordered[middle - 1] / 2.0
+ ordered[middle] / 2.0;
}
};
// Driver code
int main() {
Solution obj;
obj.addNum(1);
obj.addNum(2);
cout << obj.findMedian() << "\n";
obj.addNum(3);
cout << obj.findMedian() << "\n";
return 0;
}

Complexity Analysis

Time Complexity: O(N log N) per findMedian() query, where N is the number of values currently stored in the stream, because all N values are copied and sorted. addNum() takes O(1) time for a simple append.

Space Complexity: O(N), because the copied sorted list stores all N stream values in addition to the persistent stream storage.

Better Approach

The brute force method repeats a full sort for every query. A sorted list keeps earlier ordering work and places each arriving value directly into the correct position.

The middle value then stays available without another sort. Array insertion can shift many larger values, so updates remain linear even though median queries become constant time.

Algorithm

  • Begin with an empty sorted list, preserving ascending order from the first insertion onward.

  • Locate the first position containing a value not smaller than the new value, keeping duplicate placement valid and predictable.

  • Insert the new value at the located position because all earlier entries are smaller and all later entries are at least as large.

  • Shift later array entries because array insertion must open one position while preserving the complete sorted order.

  • Return index n / 2 for an odd size because the sorted list exposes the single middle value immediately.

  • Average indices n / 2 - 1 and n / 2 for an even size, avoiding another traversal or sorting pass.

Dry Run

Median from Data Stream - Better Approach

Median from Data Stream - Better Approach

Solution

#include <bits/stdc++.h>
using namespace std;
class Solution {
private:
vector<int> ordered;
public:
// Create empty sorted stream storage.
Solution() {
}
// Add one value while preserving sorted order.
void addNum(int value) {
// Find the first valid insertion position.
auto position = lower_bound(
ordered.begin(), ordered.end(), value
);
// Insert before the first value not smaller.
ordered.insert(position, value);
}
// Return the median from the sorted stream.
double findMedian() {
int size = ordered.size();
int middle = size / 2;
// Odd size has one value at the middle index.
if (size % 2 == 1) {
return ordered[middle];
}
// Half-value addition avoids integer overflow.
return ordered[middle - 1] / 2.0
+ ordered[middle] / 2.0;
}
};
// Driver code
int main() {
Solution obj;
obj.addNum(1);
obj.addNum(2);
cout << obj.findMedian() << "\n";
obj.addNum(3);
cout << obj.findMedian() << "\n";
return 0;
}

Complexity Analysis

Time Complexity: O(N) per addNum() operation, where N is the number of values currently stored in the stream, because insertion may shift up to N values. findMedian() takes O(1) time by accessing one or two direct indices.

Space Complexity: O(N), because the sorted list stores all N stream values.

Optimal Approach

The sorted-list method maintains more order than a median actually needs. We only need to separate the values into a smaller half and a larger half. A max-heap stores the smaller half, while a min-heap stores the larger half, so the two values closest to the median are always available at the heap roots.

The approach maintains two important invariants after every insertion:

  • Every value in the max-heap is less than or equal to every value in the min-heap.

  • The max-heap has either the same number of values as the min-heap or exactly one extra value.

These invariants keep the median at the boundary between the two heaps.

Algorithm

  • Create a max-heap for the smaller half and a min-heap for the larger half.

  • For every new value:

    1. Insert it into the max-heap first.

    2. Move the max-heap root to the min-heap so that every value in the smaller half remains less than or equal to every value in the larger half.

    3. If the min-heap becomes larger, move its root back to the max-heap so the max-heap has either the same size or one extra value.

  • After every insertion, maintain the invariant that:

    • all max-heap values are <= all min-heap values, and

    • maxHeap.size() is either equal to minHeap.size() or exactly one greater.

  • For an odd number of values, return the max-heap root because it contains the single middle value.

  • For an even number of values, return the average of both heap roots because they are the two middle values.

Dry Run

Diagram 1
1 / 2

Diagram 1

Solution

#include <bits/stdc++.h>
using namespace std;
class Solution {
private:
priority_queue<int> lower;
priority_queue<int, vector<int>, greater<int>> upper;
public:
// Create two empty stream halves.
Solution() {
}
// Add one value and restore both heap rules.
void addNum(int value) {
// Treat the new value as a lower-half candidate.
lower.push(value);
// Move the largest lower value across the boundary.
upper.push(lower.top());
lower.pop();
// Extra upper size must return one middle value.
if (upper.size() > lower.size()) {
lower.push(upper.top());
upper.pop();
}
}
// Return the median from the heap boundary.
double findMedian() {
// Equal sizes place two values at the middle.
if (lower.size() == upper.size()) {
// Half-value addition avoids integer overflow.
return lower.top() / 2.0 + upper.top() / 2.0;
}
// Extra lower value becomes the single middle.
return lower.top();
}
};
// Driver code
int main() {
Solution obj;
obj.addNum(1);
obj.addNum(2);
cout << obj.findMedian() << "\n";
obj.addNum(3);
cout << obj.findMedian() << "\n";
return 0;
}

Complexity Analysis

Time Complexity: O(log N) per addNum() operation, where N is the number of values currently stored in the stream, because each call performs a constant number of heap operations. findMedian() takes O(1) time by reading at most two heap roots.

Space Complexity: O(N), because the two heaps together store all N stream values.

Interview follow-up Questions

A single heap exposes only one extreme. Two heaps expose the largest lower-half value and the smallest upper-half value, exactly matching the one or two positions needed for a median.

Heap

Read Similar Blogs

Comments0