Stack that does Push, Pop and Median in O(1) Here are two implementation of a stack that does Push, Pop and Median in O(1). Invite comments on complexity and betterment.


  Option - 1


public class MedianStack {

    private Deque&lt;Integer&gt; stack;
    private PriorityQueue&lt;Integer&gt; minHeap;
    private PriorityQueue&lt;Integer&gt; maxHeap;

    public MedianStack(){
        stack = new ArrayDeque&lt;&gt;();
        minHeap = new PriorityQueue&lt;&gt;();
        maxHeap = new PriorityQueue&lt;&gt;(Collections.reverseOrder());
    }
    public void push(int item){
        stack.offerFirst(item);
        minHeap.add(item);
        maxHeap.add(item);
    }
    public int pop(){
        int item = stack.poll();
        minHeap.remove(item);
        maxHeap.remove(item);
        return item;
    }

    public int median(){
        return (minHeap.peek() + maxHeap.peek())/2;

    }
}



  Option - 2


Modified and added after realizing the first implementation was not correct.

public class MedianStack {

    private Deque&lt;Integer&gt; stack;
    private MedianHeap medianHeap;


    public MedianStack(){
        stack = new ArrayDeque&lt;&gt;();
        medianHeap = new MedianHeap();

    }
    public void push(int item){
        stack.offerFirst(item);
        medianHeap.add(item);

    }
    public int pop(){
        int item = stack.poll();
        medianHeap.remove(item);
        return item;
    }

    public int median(){
        return medianHeap.median();
    }
}


Median Heap Implementation

public class MedianHeap {

    private PriorityQueue&lt;Integer&gt; maxHeap;
    private PriorityQueue&lt;Integer&gt; minHeap;

    public MedianHeap(){
        maxHeap = new PriorityQueue&lt;&gt;(Collections.reverseOrder());
        minHeap = new PriorityQueue&lt;&gt;();
    }

    public void add(int item){
        if(isEmpty()) minHeap.add(item);
        else{
            if(Integer.compare(item, median()) &gt;= 0) minHeap.add(item);
            else maxHeap.add(item);
        }
        balance();
    }

    public boolean isEmpty(){
        return maxHeap.size() == 0 &amp;&amp; minHeap.size() == 0;
    }
    public void balance(){
        if(Math.abs(maxHeap.size() - minHeap.size()) &gt; 1){
            if(minHeap.size() &gt; maxHeap.size()){
                maxHeap.add(minHeap.poll());
            }else minHeap.add(maxHeap.poll());

        }
    }
    public void remove(int item){
        if(!minHeap.contains(item) || !maxHeap.contains(item)) throw new IllegalStateException("Illegal item removal");
        if(minHeap.contains(item)) minHeap.remove(item);
        else maxHeap.remove(item);
        balance();
    }

    public int median(){
        if (minHeap.size() == maxHeap.size()) return (maxHeap.peek() + minHeap.peek())/2;
        else if(minHeap.size() &gt; maxHeap.size()) return minHeap.peek();
        else return maxHeap.peek();
    }
}
minHeap.add(item);
        maxHeap.add(item);



Adding to a PriorityQueue is not \$O(1)\$.  It is \$O(\log{n})\$.  

Same thing for removal.  


    public int median(){
        return (minHeap.peek() + maxHeap.peek())/2;

    }



This is not a median.  The smallest and largest elements are totally irrelevant to a median.  In 4, 5, 6, then 5 is the median.  In -4000, 5, 6000, then 5 is still the median.  The only time you take an average to get a median is if there are an even number of elements in the set.  E.g. in 2, 4, 6, 8, then \$(4+6)/2 = 5\$ is the median.  Note that it is only the middle items in the set that matter.