Leetcode Max Stack

My solution to the question works well except for the pop max function And I supposed to retrieve the maximum element on the stock and remove it. It feels against this test case. How do I fix it?

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.

TestCase:

Input: ["MaxStack","push","push","popMax","peekMax"]
[,[5],[1],,]

Output: [null,null,null,5,-1]

Expected: [null,null,null,5,1]

class MaxStack {
var stack = [Int]()
    var maxStack = [Int]()
    /** initialize your data structure here. */
    init() {}
    
    func push(_ x: Int) {
        stack.append(x)
        if maxStack.isEmpty || x >= maxStack.last! {
            maxStack.append(x)
        }
    }
    
    func pop() -> Int {
        if let last = stack.popLast() {
            if last == maxStack.last! {
                maxStack.removeLast()
            }
            return last
        }
        return -1
    }
    
    func top() -> Int {
        if let last = stack.last {
            return last
        }
        return -1
    }
    
    func peekMax() -> Int {
        if let last = maxStack.last {
            return last
        }
        return -1
    }
    
    func popMax() -> Int {
        if let last = maxStack.popLast() {
            var max = last
            while let last = stack.popLast(), last != max {
                if last > max {
                    max = last
                }
            }
            return max
        }
        return -1
    }
}

Think about what popMax should do to stack.

It's suppose to return the max stack value and remove it. How do I resolve it ?

That's not quite the right level of abstraction. Think of each variable as a state, and that each function transitions one valid state to another valid state.

In this case, stack would be the actual stack, and maxStack would be the list of maximum values up until that point. This means that particularly, if

stack := [ 1, 2, 1, 5, 4 ]

then it follows that

maxStack := [ 1, 2, 5 ]

So when I told you to think about what the maxStack is meant to be, I mean that if you have the stack as shown above, and try to popMax, you would think about how to manually do it. Particularly,

  • You'd have to find the max element, which from maxStack, you know is 5
  • You would have to pop 5 from both stacks appropriately,
  • This means that it becomes
    stack := [ 1, 2, 1, 4 ]
    maxStack := [ 1, 2, 4 ]
    

And that's not quite how your implementation of popMax would end up if you start at the original state.

Now, you could try to figure out how to make it work with stack and maxStack. You could also try other structures as well and see if it makes maintaining these states easier. What if you don't use maxStack? What if you only store the indices to the max element instead? Etc.

One easy fix could be use some duplication.
Think about this case:
stack := [ 1, 2, 1, 5, 4 ]
maxStack := [ 1, 2, 2, 5, 5 ]
So, every time you push a new member to stack, you could either append a new max/member, or append a new member while appending an existing max member to maxStack.
And, every time you pop, you could just remove last elements for both array.

That is fine for poping, but what about max poping? By that convention, it'd have to become:

stack := [ 1, 2, 1, 4 ]
maxStack := [ 1, 2, 2, 4 ]

This doesn't look easy to do.

But yeah, the main approach when solving problems like this is to figure out the best way to organize states (stack, maxStack, etc.) so that the functions (push, pop, top, peekMax, popMax) are easy to implement.

I'll leave ppl here to have some fun with it.

1 Like