Problem

We need to design a stack called MaxStack that has the typical push, pop and top functionalities of a stack but also includes peekMax and popMax functions.

• push(x) adds element x onto stack.
• pop() removes and returns the element on top of the stack.
• top() returns the element on the top without removing it from the stack.
• peekMax() returns the maximum element in the stack without removing it from the stack.
• popMax() removes the top-most maximum element from the stack and returns it.

If popMax() method finds more than one maximum elements, it should only remove the top-most one.

Let's walk through an example to better understand this:

1
2
3MaxStack stack = new MaxStack();
4stack.push(5);    // stack: [5]
5stack.push(1);    // stack: [5,1]
6stack.push(5);    // stack: [5,1,5]
7stack.top();      // returns 5
8stack.popMax();   // returns 5, stack: [5,1]
9stack.top();      // returns 1
10stack.peekMax();  // returns 5
11stack.pop();      // returns 1, stack: [5]
12stack.top();      // returns 5

Approach

We will use two data structures in the solution. A LinkedList and a HashMap. The reasoning behind using a linked list for the stack is that, it allows to do insertions and deletions at both ends in constant O(1) time which is ideal for implementing a stack. The HashMap's key will hold the actual elements in the stack and its value will be a list that stores iterators to the linked list where the same element is found.

Anytime we push an element onto the stack, we add it to the head of the linked list and also map it into the HashMap. The HashMap allows us to look up the maximum value in constant time using the peekMax() method. The popMax() also uses the HashMap to locate and remove the maximum value from the linked list in constant time. When we pop an element off the stack, we also check if it exists in the HashMap and remove it if necessary.

Python Solution

1
2python
3from collections import defaultdict
4import heapq
5
6class MaxStack:
7
8    def __init__(self):
9        self.maxHeap = [] # max heap to store elements by their val and index
10        self.stack = [] # stack to store the elements
11        self.toPop = defaultdict(int) # hash map to track the deleted elements from max heap
12
13    def push(self, x: int) -> None:
14        heapq.heappush(self.maxHeap, (-x, -len(self.stack))) # push into max heap
15        self.stack.append(x) # push into stack
16
17    def pop(self) -> int:
18        self.toPop[self.stack[-1]] += 1 # mark as to be removed from max heap
19        return self.stack.pop() # pop from stack
20
21    def top(self) -> int:
22        return self.stack[-1] # return top of stack
23
24    def peekMax(self) -> int:
25        return -self.maxHeap[0][0] # return max of heap
26
27    def popMax(self) -> int:
28        val, _ = heapq.heappop(self.maxHeap) # pop max from heap
29        self.toPop[-val] -= 1 # decrease counter in toPop
30        while self.toPop[self.stack[-1]]: # keep popping from stack until it's empty or top of stack is not marked for deletion
31            self.toPop[self.stack[-1]] -= 1
32            self.stack.pop()
33        return -val # return the max value

JavaScript Solution

Here's a JavaScript solution that follows a similar approach as the Python solution, using the array data structure which has push, pop, peek functionality built-in. Javascript does not support the max heap, hence the peekMax and popMax operations are not in constant time but linear.

1
2javascript
3class MaxStack {
4    constructor() {
5        this.stack = [];
6    }
7
8    push(x) {
9        this.stack.push(x);
10    }
11
12    pop() {
13        return this.stack.pop();
14    }
15
16    top() {
17        return this.stack[this.stack.length - 1];
18    }
19
20    peekMax() {
21        return Math.max(...this.stack);
22    }
23
24    popMax() {
25        let max = Math.max(...this.stack);
26        for (let i = this.stack.length - 1; i >= 0; i--) {
27            if (this.stack[i] === max) {
28                this.stack.splice(i, 1);
29                break;
30            }
31        }
32        return max;
33    }
34}

Java Solution

The Java version follows the same concept, using a LinkedList for the stack and TreeMap to to keep track of maximum elements where keys are the actual elements of the stack and values are counts of the stack elements.

1
2java
3class MaxStack {
4
5    private TreeMap<Integer, List<Node>> treeMap;
6    private DoubleLinkedList dll;
7
8    /** initialize your data structure here. */
9    public MaxStack() {
10        treeMap = new TreeMap<>();
11        dll = new DoubleLinkedList();
12    }
13
14    public void push(int x) {
15        Node node = dll.add(x);
16        if(!treeMap.containsKey(x))
17            treeMap.put(x, new ArrayList<>());
18        treeMap.get(x).add(node);
19    }
20
21    public int pop() {
22        int val = dll.pop();
23        List<Node> l = treeMap.get(val);
24        l.remove(l.size()-1);
25        if(l.isEmpty())
26            treeMap.remove(val);
27        return val;
28    }
29
30    public int top() {
31        return dll.peek();
32    }
33
34    public int peekMax() {
35        return treeMap.lastKey();
36    }
37
38    public int popMax() {
39        int max = peekMax();
40        List<Node> nodes = treeMap.get(max);
41        Node node = nodes.remove(nodes.size()-1);
42        dll.unlink(node);
43        if(nodes.isEmpty())
44            treeMap.remove(max);
45        return max;
46    }
47}
48class DoubleLinkedList {
49    Node head, tail;
50    public DoubleLinkedList() {
51        head = new Node(0);
52        tail = new Node(0);
53        head.next = tail;
54        tail.prev = head;
55    }
56    public Node add(int val) {
57        Node node = new Node(val);
58        Node last = tail.prev;
59        last.next = node;
60        node.prev = last;
61        node.next = tail;
62        tail.prev = node;
63        return node;
64    }
65    public int pop() {
66        return unlink(tail.prev).val;
67    }
68    public Node unlink(Node node) {
69        Node nextNode = node.next;
70        Node prevNode = node.prev;
71        nextNode.prev = prevNode;
72        prevNode.next = nextNode;
73        return node;
74    }
75    public int peek() {
76        return tail.prev.val;
77    }
78}
79class Node {
80    Node prev, next;
81    int val;
82    public Node(int val) {
83        this.val = val;
84    }
85}

In conclusion, a MaxStack can be developed by using a combination of built-in data structures and we can add extra functionalities which can help us in many real-time scenarios.