# 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<Integer> stack;
private PriorityQueue<Integer> minHeap;
private PriorityQueue<Integer> maxHeap;

public MedianStack(){
stack = new ArrayDeque<>();
minHeap = new PriorityQueue<>();
maxHeap = new PriorityQueue<>(Collections.reverseOrder());
}
public void push(int item){
stack.offerFirst(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<Integer> stack;
private MedianHeap medianHeap;

public MedianStack(){
stack = new ArrayDeque<>();
medianHeap = new MedianHeap();

}
public void push(int item){
stack.offerFirst(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<Integer> maxHeap;
private PriorityQueue<Integer> minHeap;

public MedianHeap(){
maxHeap = new PriorityQueue<>(Collections.reverseOrder());
minHeap = new PriorityQueue<>();
}

else{
}
balance();
}

public boolean isEmpty(){
return maxHeap.size() == 0 && minHeap.size() == 0;
}
public void balance(){
if(Math.abs(maxHeap.size() - minHeap.size()) > 1){
if(minHeap.size() > maxHeap.size()){

}
}
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() > maxHeap.size()) return minHeap.peek();
else return maxHeap.peek();
}
}


        minHeap.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.

• Yup, I realise my implementation is wrong. – Clockwork May 30 '16 at 23:59
• Updated post, added option 2 which is a correct implementation. – Clockwork May 31 '16 at 0:47
• @Clockwork Rather than updating this post, you should ask a new question with the new code. Note that the new code is not $O(1)$ for push and pop either. – mdfst13 May 31 '16 at 1:20