Find Median from Data Stream LeetCode solution in Java

1
class MedianFinder {
2

3
  PriorityQueue maxHeap;
4
  PriorityQueue minHeap;
5

6
  public MedianFinder() {
7
    maxHeap = new PriorityQueue<Integer>((a, b) -> b - a);
8
    minHeap = new PriorityQueue<Integer>();
9
  }
10

11
  public void addNum(int num) {
12

13
    // Pushing
14
    if (maxHeap.isEmpty() || ((int) maxHeap.peek() > num)) {
15
      maxHeap.offer(num);
16
    } else {
17
      minHeap.offer(num);
18
    }
19

20
    // Balancing
21
    if (maxHeap.size() > minHeap.size() + 1) {
22
      minHeap.offer(maxHeap.peek());
23
      maxHeap.poll();
24
    } else if (minHeap.size() > maxHeap.size() + 1) {
25
      maxHeap.offer(minHeap.peek());
26
      minHeap.poll();
27
    }
28
  }
29

30
  public double findMedian() {
31

32
    // Evaluating Median
33
    if (maxHeap.size() == minHeap.size()) { // Even Number
34
      return ((int) maxHeap.peek() + (int) minHeap.peek()) / 2.0;
35
    } else { // Odd Number
36
      if (maxHeap.size() > minHeap.size()) {
37
        return (int) maxHeap.peek() + 0.0;
38
      } else { // minHeap.size() > maxHeap.size()
39
        return (int) minHeap.peek() + 0.0;
40
      }
41
    }
42
  }
43
}