0
\$\begingroup\$

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);
        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<Integer> stack;
    private MedianHeap medianHeap;


    public MedianStack(){
        stack = new ArrayDeque<>();
        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<Integer> maxHeap;
    private PriorityQueue<Integer> minHeap;

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

    public void add(int item){
        if(isEmpty()) minHeap.add(item);
        else{
            if(Integer.compare(item, median()) >= 0) minHeap.add(item);
            else maxHeap.add(item);
        }
        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()){
                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() > maxHeap.size()) return minHeap.peek();
        else return maxHeap.peek();
    }
}
\$\endgroup\$
6
\$\begingroup\$
        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.

\$\endgroup\$
  • \$\begingroup\$ Yup, I realise my implementation is wrong. \$\endgroup\$ – Clockwork May 30 '16 at 23:59
  • \$\begingroup\$ Updated post, added option 2 which is a correct implementation. \$\endgroup\$ – Clockwork May 31 '16 at 0:47
  • \$\begingroup\$ @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. \$\endgroup\$ – mdfst13 May 31 '16 at 1:20

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.