LeetCode/FindMedianFromDataStream.java at master · yanguango/LeetCode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
class MedianFinder {
    private PriorityQueue<Integer> leftHeap;
    private PriorityQueue<Integer> rightHeap;
    private double median;

    public MedianFinder() {
        leftHeap = new PriorityQueue(16, Collections.reverseOrder());
        rightHeap = new PriorityQueue();
    }

    // Adds a number into the data structure.
    public void addNum(int num) {
        int leftSize = leftHeap.size();
        int rightSize = rightHeap.size();
        if (leftSize < rightSize) {
            if (num < median) {
                leftHeap.offer(num);
            } else {
                leftHeap.offer(rightHeap.poll());
                rightHeap.offer(num);
            }
            median = (leftHeap.peek() + rightHeap.peek()) / 2.0;
        } else if (leftSize == rightSize) {
            if (num < median) {
                leftHeap.offer(num);
                median = leftHeap.peek();
            } else {
                rightHeap.offer(num);
                median = rightHeap.peek();
            }
        } else {
            if (num < median) {
                rightHeap.offer(leftHeap.poll());
                leftHeap.offer(num);
            } else {
                rightHeap.offer(num);
            }
            median = (leftHeap.peek() + rightHeap.peek()) / 2.0;
        }
    }

    // Returns the median of current data stream
    public double findMedian() {
        return median;
    }
};

// Your MedianFinder object will be instantiated and called as such:
// MedianFinder mf = new MedianFinder();
// mf.addNum(1);
// mf.findMedian();