716. Max Stack

Design a max stack data structure that supports the stack operations and supports finding the stack's maximum element.

Implement the MaxStack class:

• MaxStack() Initializes the stack object.
• void push(int x) Pushes element x onto the stack.
• int pop() Removes the element on top of the stack and returns it.
• int top() Gets the element on the top of the stack without removing it.
• int peekMax() Retrieves the maximum element in the stack without removing it.
• int popMax() Retrieves the maximum element in the stack and removes it. If there is more than one maximum element, only remove the top-most one.

 

Example 1:

Input
["MaxStack", "push", "push", "push", "top", "popMax", "top", "peekMax", "pop", "top"]
[[], [5], [1], [5], [], [], [], [], [], []]
Output
[null, null, null, null, 5, 5, 1, 5, 1, 5]

Explanation
MaxStack stk = new MaxStack();
stk.push(5);   // [5] the top of the stack and the maximum number is 5.
stk.push(1);   // [5, 1] the top of the stack is 1, but the maximum is 5.
stk.push(5);   // [5, 1, 5] the top of the stack is 5, which is also the maximum, because it is the top most one.
stk.top();     // return 5, [5, 1, 5] the stack did not change.
stk.popMax();  // return 5, [5, 1] the stack is changed now, and the top is different from the max.
stk.top();     // return 1, [5, 1] the stack did not change.
stk.peekMax(); // return 5, [5, 1] the stack did not change.
stk.pop();     // return 1, [5] the top of the stack and the max element is now 5.
stk.top();     // return 5, [5] the stack did not change.

 

Constraints:

• -107 <= x <= 107
• At most 104 calls will be made to push, pop, top, peekMax, and popMax.
• There will be at least one element in the stack when pop, top, peekMax, or popMax is called.

 

Follow up: Could you come up with a solution that supports O(1) for each top call and O(logn) for each other call?  

My answer: stack + heap

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59

class MaxStack { // alternative 1: use stack for max + normal stack, which always pushes max value everytime push(), // it also pops when normal pop()/popMax() -> O(N)/O(N) // alternative 2: Double Linked List + TreeMap LinkedList<Integer> stack; PriorityQueue<Integer> heap; /** initialize your data structure here. */ public MaxStack() { stack = new LinkedList<Integer>(); heap = new PriorityQueue<Integer>((a, b) -> {return b.compareTo(a);}); } public void push(int x) { stack.addLast(x); heap.add(x); } public int pop() { int removed = stack.removeLast(); heap.remove(removed); return removed; } public int top() { return stack.peekLast(); } public int peekMax() { return heap.peek(); } public int popMax() { System.out.println(stack); int removed = heap.poll(); int removeIndex = -1; for (int i = stack.size() - 1 ; i >= 0 ; i --) { if (stack.get(i) == removed) { removeIndex = i; break; } } stack.remove(removeIndex); System.out.println(stack); return removed; } } /** * Your MaxStack object will be instantiated and called as such: * MaxStack obj = new MaxStack(); * obj.push(x); * int param_2 = obj.pop(); * int param_3 = obj.top(); * int param_4 = obj.peekMax(); * int param_5 = obj.popMax(); */