Update 295_Find_Median_from_Data_Stream.java

Author: seanprashad

Two Heaps/295_Find_Median_from_Data_Stream.java

@@ -1,21 +1,33 @@
 class MedianFinder {
-    PriorityQueue<Integer> lower, upper;
+    PriorityQueue<Integer> p1, p2;
 
+    /** initialize your data structure here. */
     public MedianFinder() {
-        lower = new PriorityQueue<>(Comparator.reverseOrder());
-        upper = new PriorityQueue<>();
+        p1 = new PriorityQueue<>(Collections.reverseOrder());
+        p2 = new PriorityQueue<>();
     }
 
     public void addNum(int num) {
-        lower.offer(num);
-        upper.offer(lower.poll());
+        if (p1.isEmpty() || p1.peek() > num) {
+            p1.offer(num);
+        } else {
+            p2.offer(num);
+        }
 
-        if (lower.size() < upper.size()) {
-            lower.offer(upper.poll());
+        if (p1.size() > p2.size() + 1) {
+            p2.offer(p1.poll());
+        } else if (p2.size() > p1.size() + 1) {
+            p1.offer(p2.poll());
         }
     }
 
     public double findMedian() {
-        return lower.size() == upper.size() ? (double) (lower.peek() + upper.peek()) / 2 : lower.peek();
+        if (p1.size() == p2.size()) {
+            return (double) (p1.peek() + p2.peek()) / 2.0;
+        } else if (p1.size() > p2.size()) {
+            return (double) p1.peek();
+        } else {
+            return (double) p2.peek();
+        }
     }
 }