Two heap to save the smaller half and larger half number respectively.
We need to keep heap as balanced while insertion, i.e., the different of two heaps' size must be <= 1.
If we find the insertion will make the different larger than 1, we move the head number of a heap to the other, and then insert the number into the heap so that heap are kept balanced.
findMeadian is obvious. Only need to consider the top numbers of two heap.
public class MedianFinder { // a min heap to save the larger half numbers private PriorityQueue<Integer> minHeap; // a max heap to save the smaller half numbers private PriorityQueue<Integer> maxHeap; /** initialize your data structure here. */ public MedianFinder() { minHeap = new PriorityQueue<Integer>(); maxHeap = new PriorityQueue<Integer>( new Comparator<Integer>() { public int compare(Integer a, Integer b) { return b - a; } }); } public void addNum(int num) { // the size of two heaps are same if (minHeap.size() == 0 && maxHeap.size() == 0) { minHeap.offer(num); } else if (minHeap.size() > maxHeap.size()){ // heap's size is different, keep balance of two heaps if (num > minHeap.peek()) { // move the top of the minHeap to the maxHeap maxHeap.offer(minHeap.poll()); // add num to the minHeap minHeap.offer(num); } else { maxHeap.offer(num); } } else if (minHeap.size() < maxHeap.size()){ if (num < maxHeap.peek()) { minHeap.offer(maxHeap.poll()); maxHeap.offer(num); } else { minHeap.offer(num); } } else { if (num > minHeap.peek()) { minHeap.offer(num); } else { maxHeap.offer(num); } } } public double findMedian() { if (minHeap.size() == maxHeap.size()) { return (minHeap.peek()+maxHeap.peek())/2.0; } else if (minHeap.size() > maxHeap.size()) { return (double) minHeap.peek(); } return (double) maxHeap.peek(); } }