Find Median from Data Stream
HardPrompt
The median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value, and the median is the mean of the two middle values.
- For example, for
arr = [2,3,4], the median is3. - For example, for
arr = [2,3], the median is(2 + 3) / 2 = 2.5.
Implement the MedianFinder class:
MedianFinder()initializes theMedianFinderobject.void addNum(int num)adds the integernumfrom the data stream to the data structure.double findMedian()returns the median of all elements so far. Answers within10-5of the actual answer will be accepted.
Example 1:
Input
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]
Output
[null, null, null, 1.5, null, 2.0]
Explanation
MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1); // arr = [1]
medianFinder.addNum(2); // arr = [1, 2]
medianFinder.findMedian(); // return 1.5 (i.e., (1 + 2) / 2)
medianFinder.addNum(3); // arr[1, 2, 3]
medianFinder.findMedian(); // return 2.0
Constraints:
-105 <= num <= 105- There will be at least one element in the data structure before calling
findMedian. - At most
5 * 104calls will be made toaddNumandfindMedian.
Follow up:
- If all integer numbers from the stream are in the range
[0, 100], how would you optimize your solution? - If
99%of all integer numbers from the stream are in the range[0, 100], how would you optimize your solution?
Approaches
3 approaches with complexity analysis and trade-offs.
This is the most straightforward and intuitive approach. We store all the numbers that have been added to the stream in a dynamic array (or a list). When addNum is called, we simply append the new number to our list. When findMedian is called, we first sort the list and then find the middle element(s) to calculate the median. While simple, this method is highly inefficient for the findMedian operation, as sorting takes O(N log N) time, where N is the current number of elements in the stream.
Algorithm
- Data Structure: Use a simple list (like
ArrayListin Java) to store the numbers from the data stream. addNum(int num)Operation:- Simply append the new number
numto the end of the list.
- Simply append the new number
findMedian()Operation:- Sort the entire list in non-decreasing order.
- Determine the size of the list,
n. - If
nis odd, the median is the element at the middle index,n / 2. - If
nis even, the median is the average of the two middle elements at indicesn / 2 - 1andn / 2.
Walkthrough
In this approach, we use an ArrayList to keep all the numbers. The addNum method is very fast, as it just involves adding an element to the end of the list, which is an O(1) amortized operation. The main drawback is the findMedian method. To find the median, we must first have the numbers in sorted order. Therefore, every call to findMedian requires sorting the entire list. After sorting, we can easily access the middle element(s) in O(1) time to compute the median.
import java.util.ArrayList;import java.util.Collections;import java.util.List; class MedianFinder { private List<Integer> store; /** initialize your data structure here. */ public MedianFinder() { store = new ArrayList<>(); } public void addNum(int num) { store.add(num); } public double findMedian() { Collections.sort(store); int n = store.size(); if (n % 2 != 0) { // Odd number of elements return (double) store.get(n / 2); } else { // Even number of elements int mid1 = store.get(n / 2 - 1); int mid2 = store.get(n / 2); return (double) (mid1 + mid2) / 2.0; } }}Complexity
Time
`addNum(num)`: O(1) amortized time. `findMedian()`: O(N log N) time due to sorting the list of N elements.
Space
O(N), where N is the number of elements added to the stream. We need to store all the numbers.
Trade-offs
Pros
Very simple to understand and implement.
The
addNumoperation is fast (amortized O(1)).
Cons
findMedianoperation is very inefficient due to sorting the entire list every time it's called.This approach will be too slow and likely time out for large numbers of calls, as specified in the problem constraints.
To improve upon the previous approach's slow findMedian operation, we can try to keep the list sorted at all times. This way, finding the median becomes a trivial O(1) operation. However, the trade-off is that the addNum operation becomes more expensive. To add a new number, we first need to find its correct position in the sorted list (which can be done efficiently with binary search) and then insert it, which takes linear time as it may require shifting many elements.
Algorithm
- Data Structure: Use a list that is always kept in sorted order.
addNum(int num)Operation:- Use binary search to find the correct position for the new number
numin the sorted list. This takes O(log N) time. - Insert
numat that position. This requires shifting all subsequent elements, which takes O(N) time.
- Use binary search to find the correct position for the new number
findMedian()Operation:- Since the list is always sorted, access the middle element(s) directly.
- If the size
nis odd, return the element at indexn / 2. - If
nis even, return the average of elements atn / 2 - 1andn / 2.
Walkthrough
This method prioritizes a fast findMedian by maintaining a sorted list. When a new number arrives via addNum, we use binary search to locate the index where the new number should be inserted to preserve the sorted order. The List.add(index, element) method in Java's ArrayList then inserts the element, but this operation has a time complexity of O(N) in the worst case because it needs to shift all elements from the insertion point to the end of the list. Once the number is added, findMedian is extremely fast, as it only needs to calculate the middle index and retrieve the value(s), which is an O(1) operation.
import java.util.ArrayList;import java.util.List;import java.util.Collections; class MedianFinder { private List<Integer> store; public MedianFinder() { store = new ArrayList<>(); } public void addNum(int num) { // Find the insertion point using binary search int index = Collections.binarySearch(store, num); // If not found, binarySearch returns (-(insertion point) - 1) if (index < 0) { index = -(index + 1); } store.add(index, num); } public double findMedian() { int n = store.size(); if (n % 2 != 0) { return (double) store.get(n / 2); } else { int mid1 = store.get(n / 2 - 1); int mid2 = store.get(n / 2); return (double) (mid1 + mid2) / 2.0; } }}Complexity
Time
`addNum(num)`: O(N) time, as binary search takes O(log N) but insertion takes O(N). `findMedian()`: O(1) time.
Space
O(N), where N is the number of elements. We store all numbers in the list.
Trade-offs
Pros
The
findMedianoperation is extremely fast (O(1)).
Cons
The
addNumoperation is slow (O(N)) because inserting an element into the middle of an array-based list requires shifting elements.While better than the first approach, it's still not efficient enough for the given constraints.
The most efficient solution involves a clever use of two heaps. The core idea is to maintain two halves of the number stream: a smaller half and a larger half. We use a max-heap to store the smaller half, allowing us to find its largest element (which is close to the median) in O(1). We use a min-heap for the larger half, allowing us to find its smallest element (also close to the median) in O(1). By keeping the heaps balanced in size, the median can always be calculated from the top elements of the heaps in O(1) time. Adding a new number involves a couple of heap insertions and deletions, which costs O(log N) time.
Algorithm
- Data Structure: Use two heaps (Priority Queues in Java):
- A max-heap (
small) to store the smaller half of the numbers. - A min-heap (
large) to store the larger half of the numbers.
- A max-heap (
- Invariants:
- All numbers in
smallare less than or equal to all numbers inlarge. - The heaps are balanced in size:
small.size()is either equal tolarge.size()or one greater.
- All numbers in
addNum(int num)Operation:- Add the new number
numto thesmall(max-heap). - To maintain the first invariant, move the largest element from
small(its root) tolarge. - To rebalance the sizes (the second invariant), if
largenow has more elements thansmall, move the smallest element fromlarge(its root) back tosmall.
- Add the new number
findMedian()Operation:- If the total number of elements is odd, the
smallheap will have one more element. The median is the root ofsmall. - If the total number of elements is even, the heaps will have equal size. The median is the average of the roots of
smallandlarge.
- If the total number of elements is odd, the
Walkthrough
This approach divides the numbers into two balanced halves. The small heap is a max-heap that stores the smaller half of the numbers, and large is a min-heap that stores the larger half. We maintain two key invariants: first, every number in small is less than or equal to every number in large, and second, the size of small is always equal to or one greater than the size of large.
When adding a number, we follow a simple balancing routine: add the new number to small, then move the largest element from small to large. This ensures the first invariant. Then, if large becomes bigger than small, we move the smallest element from large back to small to restore the size balance. This elegant process ensures both invariants are always met.
Finding the median is then straightforward: if the total number of elements is odd, the median is the extra element in the small heap (its root). If the count is even, the median is the average of the roots of both heaps.
import java.util.PriorityQueue;import java.util.Collections; class MedianFinder { // Max-heap for the smaller half of the numbers private PriorityQueue<Integer> small; // Min-heap for the larger half of the numbers private PriorityQueue<Integer> large; public MedianFinder() { small = new PriorityQueue<>(Collections.reverseOrder()); // Max-heap large = new PriorityQueue<>(); // Min-heap } public void addNum(int num) { // Add to the max-heap (small half) small.add(num); // Ensure all elements in small are <= elements in large // by moving the largest from small to large. if (!small.isEmpty() && !large.isEmpty() && small.peek() > large.peek()) { large.add(small.poll()); } // A simpler, more robust balancing logic is often used: // 1. Add to small: small.add(num); // 2. Move largest from small to large: large.add(small.poll()); // 3. If sizes are imbalanced, move smallest from large to small: if (large.size() > small.size()) { small.add(large.poll()); } // Balance the sizes if (small.size() > large.size() + 1) { large.add(small.poll()); } if (large.size() > small.size()) { small.add(large.poll()); } } public double findMedian() { if (small.size() > large.size()) { // Total number of elements is odd return small.peek(); } else { // Total number of elements is even return (small.peek() + large.peek()) / 2.0; } }}Follow-up Questions:
- If all numbers are in
[0, 100]: We can use an integer arraycountsof size 101 as a frequency map.addNum(num)becomes an O(1) operation (counts[num]++).findMedian()would involve iterating through this array to find the middle element(s), which takes O(K) time where K is the range (101). Since K is constant, this is an O(1) operation. The space complexity would also be O(K), i.e., O(1). - If 99% of numbers are in
[0, 100]: A hybrid approach could be used. We can use the frequency map for numbers in the[0, 100]range and separate data structures (like two heaps or sorted lists) for the outliers (numbers < 0 or > 100).findMedianwould then involve a more complex calculation to locate the median across these three groups. However, given that the two-heap solution is already very efficient (O(log N)), it might be a practical choice even in this scenario, as it gracefully handles outliers without added implementation complexity.
Complexity
Time
`addNum(num)`: O(log N) time, due to heap insertion/deletion. `findMedian()`: O(1) time, as it only involves peeking at the top of the heap(s).
Space
O(N), where N is the number of elements. All numbers are stored in the two heaps.
Trade-offs
Pros
Highly efficient with O(log N) for
addNumand O(1) forfindMedian.The ideal solution for streaming data where both adding and querying are frequent.
Scales well as the number of elements grows.
Cons
Implementation is more complex than the list-based approaches.
Requires a good understanding of heap data structures.
Solutions
Solution
public class MedianFinder { private List < int > nums ; private int curIndex ; /** initialize your data structure here. */ public MedianFinder () { nums = new List < int >(); } private int FindIndex ( int val ) { int left = 0 ; int right = nums . Count - 1 ; while ( left <= right ) { int mid = left + ( right - left ) / 2 ; if ( val > nums [ mid ]) { left = mid + 1 ; } else { right = mid - 1 ; } } return left ; } public void AddNum ( int num ) { if ( nums . Count == 0 ) { nums . Add ( num ); curIndex = 0 ; } else { curIndex = FindIndex ( num ); if ( curIndex == nums . Count ) { nums . Add ( num ); } else { nums . Insert ( curIndex , num ); } } } public double FindMedian () { if ( nums . Count % 2 == 1 ) { return ( double ) nums [ nums . Count / 2 ]; } else { if ( nums . Count == 0 ) { return 0 ; } else { return ( double ) ( nums [ nums . Count / 2 - 1 ] + nums [ nums . Count / 2 ]) / 2 ; } } } } /** * Your MedianFinder object will be instantiated and called as such: * MedianFinder obj = new MedianFinder(); * obj.AddNum(num); * double param_2 = obj.FindMedian(); */class MedianFinder { private PriorityQueue < Integer > q1 = new PriorityQueue <>(); private PriorityQueue < Integer > q2 = new PriorityQueue <>( Collections . reverseOrder ()); /** initialize your data structure here. */ public MedianFinder () { } public void addNum ( int num ) { q1 . offer ( num ); q2 . offer ( q1 . poll ()); if ( q2 . size () - q1 . size () > 1 ) { q1 . offer ( q2 . poll ()); } } public double findMedian () { if ( q2 . size () > q1 . size ()) { return q2 . peek (); } return ( q1 . peek () + q2 . peek ()) * 1.0 / 2 ; } } /** * Your MedianFinder object will be instantiated and called as such: * MedianFinder obj = new MedianFinder(); * obj.addNum(num); * double param_2 = obj.findMedian(); *//** * initialize your data structure here. */ var MedianFinder = function () { this.val = [];};/** * @param {number} num * @return {void} */ MedianFinder.prototype.addNum = function (num) { let left = 0; let right = this.val.length; while (left < right) { let mid = left + ~~((right - left) / 2); if (num > this.val[mid]) { left = mid + 1; } else { right = mid; } } this.val.splice(left, 0, num); };/** * @return {number} */ MedianFinder.prototype.findMedian = function () { let mid = ~~(this.val.length / 2); return this.val.length % 2 ? this.val[mid] : (this.val[mid - 1] + this.val[mid]) / 2;};from heapq import heappush , heappop class MedianFinder : def __init__ ( self ): # the smaller half of the list, max heap (invert min-heap) self . smallerHalf = [] # the larger half of the list, min heap self . largerHalf = [] def addNum ( self , num : int ) -> None : # trick for smaller half, use -1*val # heapq in python does NOT have comparator like in Java heappush ( self . smallerHalf , - num ) heappush ( self . largerHalf , - heappop ( self . smallerHalf )) # note: not self.smallerHalf.pop() if len ( self . smallerHalf ) < len ( self . largerHalf ): heappush ( self . smallerHalf , - heappop ( self . largerHalf )) def findMedian ( self ) -> float : if len ( self . smallerHalf ) == len ( self . largerHalf ): return ( - self . smallerHalf [ 0 ] + self . largerHalf [ 0 ]) / 2.0 else : return float ( - self . smallerHalf [ 0 ]) # Your MedianFinder object will be instantiated and called as such: # obj = MedianFinder() # obj.addNum(num) # param_2 = obj.findMedian() # follow up class MedianFinder : def __init__ ( self ): self . counts = [ 0 ] * 101 self . total = 0 def addNum ( self , num : int ) -> None : self . counts [ num ] += 1 self . total += 1 def findMedian ( self ) -> float : if self . total % 2 == 0 : # even number of elements middle1 = self . total // 2 middle2 = middle1 + 1 count = 0 i1 = i2 = 0 for i in range ( 101 ): count += self . counts [ i ] if count >= middle1 and i1 == 0 : i1 = i if count >= middle2 : i2 = i break return ( i1 + i2 ) / 2 else : # odd number of elements middle = self . total // 2 + 1 count = 0 for i in range ( 101 ): count += self . counts [ i ] if count >= middle : return i ############ class MedianFinder : def __init__ ( self ): """ initialize your data structure here. """ self . h1 = [] self . h2 = [] def addNum ( self , num : int ) -> None : heappush ( self . h1 , num ) heappush ( self . h2 , - heappop ( self . h1 )) if len ( self . h2 ) - len ( self . h1 ) > 1 : heappush ( self . h1 , - heappop ( self . h2 )) def findMedian ( self ) -> float : if len ( self . h2 ) > len ( self . h1 ): return - self . h2 [ 0 ] return ( self . h1 [ 0 ] - self . h2 [ 0 ]) / 2 # Your MedianFinder object will be instantiated and called as such: # obj = MedianFinder() # obj.addNum(num) # param_2 = obj.findMedian()class MedianFinder { public: /** initialize your data structure here. */ MedianFinder () { } void addNum ( int num ) { q1 . push ( num ); q2 . push ( q1 . top ()); q1 . pop (); if ( q2 . size () - q1 . size () > 1 ) { q1 . push ( q2 . top ()); q2 . pop (); } } double findMedian () { if ( q2 . size () > q1 . size ()) { return q2 . top (); } return ( double ) ( q1 . top () + q2 . top ()) / 2 ; } private: priority_queue < int , vector < int > , greater < int >> q1 ; priority_queue < int > q2 ; }; /** * Your MedianFinder object will be instantiated and called as such: * MedianFinder* obj = new MedianFinder(); * obj->addNum(num); * double param_2 = obj->findMedian(); */Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.