Short simple Java/C++/Python, O(log n) + O(1)


  • 179

    I keep two heaps (or priority queues):

    • Max-heap small has the smaller half of the numbers.
    • Min-heap large has the larger half of the numbers.

    This gives me direct access to the one or two middle values (they're the tops of the heaps), so getting the median takes O(1) time. And adding a number takes O(log n) time.

    Supporting both min- and max-heap is more or less cumbersome, depending on the language, so I simply negate the numbers in the heap in which I want the reverse of the default order. To prevent this from causing a bug with -231 (which negated is itself, when using 32-bit ints), I use integer types larger than 32 bits.

    Using larger integer types also prevents an overflow error when taking the mean of the two middle numbers. I think almost all solutions posted previously have that bug.

    Update: These are pretty short already, but by now I wrote even shorter ones.


    Java

    class MedianFinder {
    
        private Queue<Long> small = new PriorityQueue(),
                            large = new PriorityQueue();
    
        public void addNum(int num) {
            large.add((long) num);
            small.add(-large.poll());
            if (large.size() < small.size())
                large.add(-small.poll());
        }
    
        public double findMedian() {
            return large.size() > small.size()
                   ? large.peek()
                   : (large.peek() - small.peek()) / 2.0;
        }
    };
    

    Props to larrywang2014's solution for making me aware that I can use Queue in the declaration instead of PriorityQueue (that's all I got from him, though (just saying because I just saw he changed his previously longer addNum and it's now equivalent to mine)).


    C++

    class MedianFinder {
        priority_queue<long> small, large;
    public:
    
        void addNum(int num) {
            small.push(num);
            large.push(-small.top());
            small.pop();
            if (small.size() < large.size()) {
                small.push(-large.top());
                large.pop();
            }
        }
    
        double findMedian() {
            return small.size() > large.size()
                   ? small.top()
                   : (small.top() - large.top()) / 2.0;
        }
    };
    

    Big thanks to jianchao.li.fighter for telling me that C++'s priority_queue is a max-queue (see comments below).


    Python

    from heapq import *
    
    class MedianFinder:
    
        def __init__(self):
            self.heaps = [], []
    
        def addNum(self, num):
            small, large = self.heaps
            heappush(small, -heappushpop(large, num))
            if len(large) < len(small):
                heappush(large, -heappop(small))
    
        def findMedian(self):
            small, large = self.heaps
            if len(large) > len(small):
                return float(large[0])
            return (large[0] - small[0]) / 2.0

  • 1

    Hi, Stefan. In C++, the default priority_queue is a max heap. So I guess in order to fit with your Java and Python codes, the C++ code may require some revisions: large (instead of small) needs to use negated numbers. Do you think so? Anyway, the idea of using negated numbers is really cool :-)


  • 1

    Meh, that's normal. I just talked about it for completeness. Quoting a top-voted answer and responses:

    Daniel: The easiest way is to invert the value of the keys and use heapq. For example, turn 1000.0 into -1000.0 and 5.0 into -5.0.
    Douglas: That's kind of a kludgy solution...
    Andrew: It's also the standard solution.

    :-D


  • 0

    Hi, Stefan, I have updated my original comment now :-) Sorry for the delay.

    In fact, we need to make large a min heap and small a max heap. So in C++, large actually needs to use negated numbers. Your original code actually stores the larger half in small and the smaller half in large. I make a little change to your original C++ code as follows to fit with that design. But, how do you think?

    class MedianFinder {
        priority_queue<long> small, large;
    public:
    
        void addNum(int num) {
            large.push(-num); // adding one negative sign
            small.push(-large.top()); 
            large.pop();
            if (small.size() > large.size()) {
                large.push(-small.top());
                small.pop();
            }
        }
    
        double findMedian() {
            return small.size() < large.size()
                   ? -large.top() // adding another negative sign 
                   : -(large.top() - small.top()) / 2.0; // adding the third negative sign 
        }
    };

  • 1

    Aaaahahahahahaha... that's hilarious! You're right, in C++ the default is max. I thought it's min. So for like two minutes I was baffled why my solution still works. Then I realized that the whole thing is so symmetric that it doesn't matter. Just the names "small" and "large" are wrong, but names don't matter except for the reader. Oh gawd... now I have to rewrite that code and text even though it works anyway. And check what Java is actually doing. Sigh...


  • 0

    Haha, the symmetry is so beautiful that the rewrite involves just adding three negatives signs :-) I check the above rewrite and it actually stores the larger/smaller half in large/small.


  • 0

    Yeah, that works. But I don't want to add minus signs, as I had already intentionally minimized them :-). So now I instead changed the names (and the explanation).


  • 0
    Y

    The answer is awesome!!!
    I like this answer!


  • 3
    N

    In C++, there is another way to represent large. In this way, we needn't use negatives signs. I will optimize my code later.

    priority_queue<int, greater<int>> large; // #include <functional> greater
    
    class MedianFinder {
    public:
    
    	// Adds a number into the data structure.
    	void addNum(int num) {
    		if (left.empty()) {
    			left.push(num);
    			return;
    		}
    		int size1 = left.size();
    		int size2 = right.size();
    		if (size1 == size2) {
    			int tmp = right.top();
    			if (num <= tmp) {
    				left.push(num);
    			}
    			else {
    				left.push(tmp);
    				right.pop();
    				right.push(num);
    			}
    		}
    		else {
    			int tmp = left.top();
    			if (tmp <= num) {
    				right.push(num);
    			}
    			else {
    				left.pop();
    				left.push(num);
    				right.push(tmp);
    			}
    		}
    	}
    
    	// Returns the median of current data stream
    	double findMedian() {
    		int size1 = left.size();
    		int size2 = right.size();
    		if (size1 > size2) {
    			return left.top();
    		}
    		else {
    			return (left.top() + right.top()) / 2.0;
    		}
    	}
    private:
    	priority_queue<int> left;
    	priority_queue<int, vector<int>, greater<int>> right;
    };

  • 2
    T

    Brilliant solution. For C++, you can simply use priority_queue<int, vector<int>, greater<int>> large to get a min heap, so that you don't have to do those minus stuff and make it clear, I think. (Thanks for telling me that, jianchao)


  • 0

    @TonyLic You can choose your code and press Ctrl + K. BTW, I guess the above alternative adds too many codes for Stefan:-)


  • 0

    Yes, it's a bit more code, but it would indeed be simpler/clearer (though at least for me, the negation trick isn't that bad, as I'm used to it). Really what happened is that I wrote it in Python first, which doesn't have a good alternative, and then I just translated it to C++ and Java.


  • 0
    J

    I really appreciate your ideas and the simple coding, but your method does lack efficiency for practical use due to unnecessary push & pop each time when adding a number.


  • 13
    C

    Hi @Stefan.
    In Java we can easily convert min heap into max heap and vice versa by overriging Comparator and passing it to the constuctor. something like this.....

    class MedianFinder {
       private Queue<Integer> small = new PriorityQueue(1, new Comparator<Integer>() {
    		   public int compare(Integer o1, Integer o2) {
    			   return o2 - o1; 
    		   };
       });
    	   
       private Queue<Integer> large = new PriorityQueue();
        // Adds a number into the data structure.
        public void addNum(int num) {
            large.add(num);
            small.add(large.poll());
            if (large.size() < small.size())
                large.add(small.poll());
        }
    
        // Returns the median of current data stream
        public double findMedian() {
            return large.size() > small.size()
                   ? large.peek()
                   : (large.peek() + small.peek()) / 2.0;
        }
    }

  • 0
    H

    Can we maxHeap instead of using Long and pushing negative number? Long would also cost double the memory.

    private static MaxHeapComparator maxHeapComparator = new MaxHeapComparator();
    	static class MaxHeapComparator implements Comparator<Integer> {
    
    		@Override
    		public int compare(Integer o1, Integer o2) {
    			return o2 - o1;
    		}
    		
    	}
    	
    	PriorityQueue<Integer> maxHeapOfSmallerElements = new PriorityQueue<>(16, maxHeapComparator);
    	PriorityQueue<Integer> minHeapOfLargerElements = new PriorityQueue<>();
    
    
    	 // Adds a number into the data structure.
        public void addNum(int num) {
        	if(minHeapOfLargerElements.size()>0 && num > minHeapOfLargerElements.peek()){
        		minHeapOfLargerElements.add(num);
        		if(minHeapOfLargerElements.size()>maxHeapOfSmallerElements.size()) {
        			maxHeapOfSmallerElements.add(minHeapOfLargerElements.poll());
        		}
        		
        	} else {
    	    	maxHeapOfSmallerElements.add(num); 
    	    	if(maxHeapOfSmallerElements.size() > minHeapOfLargerElements.size()+1) {
    	    		minHeapOfLargerElements.add(maxHeapOfSmallerElements.poll());
    	    	}
        	}
        }
    
        // Returns the median of current data stream
        public double findMedian() {
        	if((maxHeapOfSmallerElements.size() + minHeapOfLargerElements.size())%2!=0) {
        		return (double)maxHeapOfSmallerElements.peek();
        	}
        	
        	return (maxHeapOfSmallerElements.peek()+minHeapOfLargerElements.peek())/2.0;        
        }
    

  • 1
    E

    I think your method is more Understandable, thx.


  • 0
    R

    You can do Queue<Integer> small = new PriorityQueue<>(Comparator.reverseOrder());


  • 2
    D

    can use lambda expression

    public class MedianFinder {
    	 	private Queue<Integer> small = new PriorityQueue<>((o1,o2)->(o2-o1));
    		private Queue<Integer> large = new PriorityQueue<>(); //default is max
    	    // Adds a number into the data structure.
    	    public void addNum(int num) {
    	        large.add(num);
    	        small.add(large.poll());
    	        if(large.size()<small.size())	large.add(small.poll());
    	    }
    	
    	    // Returns the median of current data stream
    	    public double findMedian() {
    	        return large.size()>small.size()?large.peek(): (large.peek() + small.peek())/2.0;
    	    }
    	};
    

  • 0
    L

    I love you 😘😘😘


  • 1
    R

    @StefanPochmann I was asked this question in amazon interview last year. In my case the interviewer wanted me to use BST instead of inbuilt PQ. I am assuming that probably would be likely with other interviewers. This solution would work for solving this problem but from interview perspective the BST rank based solution posted in another thread is more valuable.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.