LeetCode link

Intuition

  • Using two heaps, maxHeap and minHeap, ensure the size of max always not smaller the min, so the median is the max peek(when max is one more than min) or the middle of max peek and min peek(when the sizes of them are equal).

solution

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
class MedianFinder {
    PriorityQueue<Integer> minHeap;
    PriorityQueue<Integer> maxHeap;
    /** initialize your data structure here. */
    public MedianFinder() {
        minHeap = new PriorityQueue<>();
        maxHeap = new PriorityQueue<>();
    }
    public void addNum(int num) {
        if (maxHeap.isEmpty()) {
            maxHeap.add(-num);
        } else {
            if (num <= -maxHeap.peek()) {
                maxHeap.add(-num);
                if (maxHeap.size() > minHeap.size() + 1) {
                    minHeap.add(-maxHeap.poll());
                }
            } else {
                if (maxHeap.size() > minHeap.size()) {
                    minHeap.add(num);
                } else {
                    minHeap.add(num);
                    maxHeap.add(-minHeap.poll());
                }
            }
        }
    }
    public double findMedian() {
        if (maxHeap.size() == minHeap.size()) {
            return (double)(minHeap.peek() - maxHeap.peek()) / 2.0;
        }
        return -maxHeap.peek();
    }
}

promote

The code can be simplify.

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
class MedianFinder {
    PriorityQueue<Integer> minHeap;
    PriorityQueue<Integer> maxHeap;
    /** initialize your data structure here. */
    public MedianFinder() {
        minHeap = new PriorityQueue<>();
        maxHeap = new PriorityQueue<>();
    }
    public void addNum(int num) {
        minHeap.add(num);
        maxHeap.add(-minHeap.poll());
        if (maxHeap.size() > minHeap.size() + 1) {
            minHeap.add(-maxHeap.poll());
        }
    }
    public double findMedian() {
        if (maxHeap.size() == minHeap.size()) {
            return (double)(minHeap.peek() - maxHeap.peek()) / 2.0;
        }
        return -maxHeap.peek();
    }
}