Reference: LeetCode
Difficulty: Easy

Problem

Design a max stack that supports push, pop, top, peekMax and popMax.

Note:

  • -1e7 <= x <= 1e7
  • Number of operations won’t exceed 10000.
  • The last four operations won’t be called when stack is empty.

Example:

1
2
3
4
5
6
7
8
9
10
MaxStack stack = new MaxStack();
stack.push(5);
stack.push(1);
stack.push(5);
stack.top(); -> 5
stack.popMax(); -> 5
stack.top(); -> 1
stack.peekMax(); -> 5
stack.pop(); -> 1
stack.top(); -> 5

Analysis

My first attempt (ListNode): gist Not Accepted

Failed in case: [3 5 1] (popMax 5, it does not know the max is then 3).

However, by writing the code I improved writing functions in ListNode.

Two Stacks (as in 155. Min Stack)

This solution is based on the brute-force solution in 155. Min Stack.

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
60
class MaxStack {  

Stack<Integer> stack = new Stack<>();
Stack<Integer> buffer = new Stack<>();
private int max = Integer.MIN_VALUE;

public void push(int x) {
stack.push(x);
if (x > max) {
max = x;
}
}

public int pop() { // assume not empty
int val = stack.pop();
if (val == max) {
// update max (find the new max)
int newMax = Integer.MIN_VALUE; // consider size = 1, min should be reset
// put aside
while (stack.size() > 0) {
int x = stack.pop();
newMax = Math.max(newMax, x);
buffer.push(x);
}
max = newMax;
// put it back
while (buffer.size() > 0) {
stack.push(buffer.pop());
}
}
return val;
}

public int top() { return stack.peek(); }

public int peekMax() { return max; }

public int popMax() {
int newMax = Integer.MIN_VALUE;
// find the max
while (stack.peek() != max) {
newMax = Math.max(newMax, stack.peek());
buffer.push(stack.pop());
}
// remove max
stack.pop();
// still put all to the buffer (we want to find newMax)
while (stack.size() > 0) {
newMax = Math.max(newMax, stack.peek());
buffer.push(stack.pop());
}
// put back
while (buffer.size() > 0) {
stack.push(buffer.pop());
}
int oldMax = max;
max = newMax;
return oldMax;
}
}

Time: pop() and popMax() take $O(N)$
Space: $O(N)$

Two Stacks

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
class MaxStack {

private Stack<Integer> stack = new Stack<>();
private Stack<Integer> maxStack = new Stack<>();
private Stack<Integer> buffer = new Stack<>();

public void push(int x) {
stack.push(x);
if (maxStack.size() > 0) {
maxStack.push(Math.max(x, maxStack.peek()));
} else {
maxStack.push(x);
}
}

public int pop() {
maxStack.pop();
return stack.pop();
}

public int top() { return stack.peek(); }
public int peekMax() { return maxStack.peek(); }

public int popMax() {
int max = maxStack.peek();
while (stack.peek() != max) {
buffer.push(pop()); // pop from two stacks
}
pop(); // remove max
while (buffer.size() > 0) {
push(buffer.pop()); // push into two stacks
}
return max;
}
}

Note:

The following popMax() is not correct. Why? Use stack.pop() or pop()? Or use which push()?

1
2
3
4
5
6
7
8
9
10
11
// bug!
public int popMax() {
while (stack.peek() != maxStack.peek()) {
buffer.push(stack.pop());
}
stack.pop(); // remove max
while (buffer.size() > 0) {
stack.push(buffer.pop());
}
return maxStack.pop();
}

Time: Only popMax() takes $O(N)$
Space: $O(N)$

The optimization can be made in finding the largest element by TreeMap.

Nice Intuition:


Comment