Median From Data Stream Editorial

Practice Problem Link: Median From Data Stream

Please make sure to try solving the problem yourself before looking at the editorial.

Problem Statement

The median of a finite list of numbers is the "middle" number when those numbers are sorted. If the size of the list is even, the median is the average of the two middle numbers.

You have an incoming stream of integers, you need to be able to find the median at any point in time.

Implement the MedianCalculator class:

• MedianCalculator() initializes the MedianCalculator object and is called at the beginning.
• void addNum(int num) adds the integer num from the data stream.
• float getMedian() returns the median of all elements so far.

getMedian will be called only after at least one number has been added.

Naive Approach

In the naive approach, we can naively keep storing the elements in a dynamic array, and keep sorting the array and calculating the median naively according to the given formula, each time the getMedian() function is called.

Analysis

• Time Complexity: addNum() -> O(1), getMedian() -> O(n * logn).
• Space Complexity: addNum() ->O(n), getMedian() -> O(n).

Implementation

C++ 

class MedianCalculator {
public:
    /** initialize your data structure here. */
	vector<int> elements;
    MedianCalculator() {
		elements.clear();
    }
    
    void addNum(int num) {
		elements.push_back(num);
    }
    
    float getMedian() {
		vector<int> aux = elements;
		sort(aux.begin(), aux.end());
		if(aux.size() % 2 == 0) {
			return (aux[aux.size() / 2] + aux[aux.size() / 2 - 1]) / 2.0;
		}
		else {
			return (float)aux[aux.size() / 2];
		}
    }
};

/**
 * Your MedianCalculator object will be instantiated and called as such:
 * MedianCalculator* obj = new MedianCalculator();
 * obj->addNum(num);
 * float median = obj->getMedian();
 */

Java

class MedianCalculator {
    /** initialize your data structure here. */
	ArrayList<Integer> elements;
    public MedianCalculator() {
		elements = new ArrayList<>();
    }
    
    public void addNum(int num) {
		elements.add(num);
    }
    
    public float getMedian() {
		Collections.sort(elements);
		int n = elements.size();
		if(n % 2 == 0) {
			int median1 = elements.get(n / 2 - 1);
			int median2 = elements.get(n / 2);
			return (median1 + median2 + 0.0f) / 2.0f;
		}
		else {
			return (float)elements.get(n / 2);
		}
    }
}

/**
 * Your MedianCalculator object will be instantiated and called as such:
 * MedianCalculator obj = new MedianCalculator();
 * obj.addNum(num);
 * float median = obj.getMedian();
 */

Optimal Approach

In the optimal approach, we can use two heaps, one max heap, and a min-heap. One heap will store the larger half of the array, and the other heap will store the smaller half of the array in descending order. When the current length is even, we take the top elements of both the heaps and take their average. Else, we return the top of the appropriate heap as the median. It can also be solved in an easier way in C++ using ordered_set data structure. Refer Implementation for better understanding.

Analysis

• Time Complexity: addNum() -> O(logn), getMedian() -> O(logn).
• Space Complexity: addNum() -> O(n), getMedian() -> O(n).

Implementation

C++

class MedianCalculator {
public:
    /** initialize your data structure here. */
	priority_queue<int> large;
	priority_queue<int, vector<int>, greater<int>> small;
    MedianCalculator() {
		while(!large.empty()) {
			large.pop();
		}
		while(!small.empty()) {
			small.pop();
		}
    }
    
    void addNum(int num) {
		large.push(num);
        small.push(large.top());
		large.pop();
        if(large.size() < small.size()) {
            large.push(small.top());
			small.pop();
        }
    }
    
    float getMedian() {
		return large.size() > small.size() ? large.top() : (large.top() + small.top()) / 2.0f;
    }
};

/**
 * Your MedianCalculator object will be instantiated and called as such:
 * MedianCalculator* obj = new MedianCalculator();
 * obj->addNum(num);
 * float median = obj->getMedian();
 */

Java

class MedianCalculator {
    /** initialize your data structure here. */
	PriorityQueue<Integer> small, large;
    public MedianCalculator() {
		small = new PriorityQueue<>((a,b) -> b-a);
        large = new PriorityQueue<>();
    }
    
    public void addNum(int num) {
		large.add(num);
        small.add(large.poll());
        if(large.size() < small.size()) {
            large.add(small.poll());
        }
    }
    
    public float getMedian() {
		return large.size() > small.size() ? large.peek() : (large.peek() + small.peek()) / 2.0f;
    }
}

/**
 * Your MedianCalculator object will be instantiated and called as such:
 * MedianCalculator obj = new MedianCalculator();
 * obj.addNum(num);
 * float median = obj.getMedian();
 */