295. Find Median from Data Stream

Problem Statement

Find Median from Data Stream - LeetCode

Pattern:

---

Solution

public class FindMedianinStream295 {
	static class MedianFinder {
		
		private PriorityQueue<Integer> lMax;
		private PriorityQueue<Integer> rMin;
		private double currMedian;
		
		// constructor
		public MedianFinder() {
			lMax = new PriorityQueue<>(Collections.reverseOrder());
			rMin = new PriorityQueue<>();
			currMedian = 0;
		}
		
		// add num to DS
		public void addNum (int num) {
			switch (lMax.size() - rMin.size()) {
				// l == r
				case 0 : {
					// add to right
					if(num >= currMedian) rMin.add(num);
					// add to left
					else lMax.add(num);
					break;
				}
				// l > r
				case 1 : {
					// add to right
					if(num >= currMedian) rMin.add(num);
					// add to left
					else {
						// balance heaps
						rMin.add(lMax.poll());
						// add num to left heap
						lMax.add(num);
					}
					break;
				}
				// r > l
				case -1 : {
					// add to right
					if(num >= currMedian){
						// balance heaps
						lMax.add(rMin.poll());
						// add to rMin
						rMin.add(num);
					}
					// add to left
					else lMax.add(num);
					break;
				}
			}
			// update median
			findMedian();
			//System.out.printf("num :%d , currMedian: %f , lMax: %d, rMin: %d\n", num, currMedian, lMax.peek(), rMin.peek());
		}
		
		// find median
	    public double findMedian (){
			int size = lMax.size() + rMin.size();
			// even size -> avg of heap peeks
			if((size & 1) == 0) currMedian = ((double) lMax.peek()+ (double) rMin.peek())/ (double)2;
			// larger heap peek is the median
			else if (lMax.size() > rMin.size()) currMedian = lMax.peek();
			else currMedian = rMin.peek();
			// return median
			return currMedian;
		}
	}
	
	// region MAIN
	public static void main(String[] args) {
	    MedianFinder mf = new MedianFinder();
		mf.addNum(1);
		mf.addNum(2);
		System.out.println(mf.findMedian());
		mf.addNum(3);
		System.out.println(mf.findMedian());
	}
	// endregion
}
 

Output:

1.5
2.0

Notes

• an array can be thought of as to heaps left and right. left heap is a maxHeap whose greatest element is less than the median, and right heap is a min heap whose smallest element is larger than the median.
• insert element in the left or right heap depending on if it is greater or smaller than median
  • update median
  • repeat till needed
• inserting elements is not super straightforward. the idea is to maintain heaps in such a way that all the elements in the leftMax heap are always less than the currentMedian and the elements in the rightMin heap are always greater than the currentMedian
  • So say 2 consecutive elements added are greater than the currentMedian, the median would shift towards the right heap.
  • this is why we must always make sure leftMax.size - rightMin.size <= 1
  • Love Babbar Video Explaining Solution