Design Max Stack Data Structure

Maximum Value in Stack. Stack Cautiously to remain alive for checking the max stack!

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

1. push(x) — Push element x onto stack.
2. pop() — Remove the element on top of the stack and return it.
3. top() — Get the element on the top.
4. peekMax() — Retrieve the maximum element in the stack.
5. popMax() — Retrieve the maximum element in the stack, and remove it. If you find more than one maximum elements, only remove the top-most one.

Example 1:

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

Problem: leetcode.com

Key:
First need to simulate stack behavior. A list or vector can be used to do that. However, popMax operation may be a linear operation O(N) if max element is not the top element.
Approach 1:
Use two stack, one for elements and second for max so far. PopMax operation will be linear in this case when searching max element in element stack. Rest of the operation will be O(1) i.e. constant.

class TwoStackSolution {
public: 
    TwoStackSolution() {}

    void push(int x) {
        st.push(x);
        // push the max of top and current element.
        int maxe = max(mst.empty() ? INT_MIN: mst.top(), x);
        mst.push(maxe);
        cout << "push: " << x << endl;
    }
        
    int pop() {
        mst.pop();
        int x = st.top();
        st.pop();
        cout << "pop: " << x << endl;
        return x;
    }
    
    int top() {
        cout << "peek stack: " << st.top()<<endl;
        return st.top();
    }
    
    int peekMax() {
        cout << "peek max: " << mst.top()<<endl;
        return mst.top();
    }
    
    int popMax() {
        // potentially O(N) solution, but all operation remain O(1)
        stack<int> tmp;
        int mx = mst.top();
        // pop from max stack
        mst.pop();

        // now find max in stack,pop and put rest of the back!
        int top = st.top();
        while (top != mx) {
            tmp.push(top);
            st.pop();
            top = st.top();
        }
        // top is the max. Pop from stack.
        st.pop();
        // put rest of them back.
        while (!tmp.empty()) {
            st.push(tmp.top());
            tmp.pop();
        }
        cout << "pop max: " << mx << endl;
        return mx;
    }
private:
    stack<int> st;
    stack<int> mst;
};

Approach 2:
To avoid linear operation for PopMax, we need something which can make it better than linear. It can be done using a balanced tree which takes logN for insert and delete. It means max time will be O(logN) for pop and push operations. Here is the implementation which keeps max information in balanced tree in std::map and keeps the address of elements in doubly-linkedlist. Why doubly linkedlist – because after delte operation, removing given element is O(1) if we already know the address, and address is given by map if looking for pop max.

class MaxStack {
public:
    /** initialize your data structure here. */
    MaxStack() {}
    
    void push(int x) {
        l.insert(l.begin(), x);
        m[x].push_back(l.begin());
    }
    
    int pop() {
        int r = *l.begin();
        m[r].pop_back();
        if (m[r].empty())
            m.erase(r);
        l.erase(l.begin());
        return r;
    }
    
    int top() {
        return *l.begin();
    }
    
    int peekMax() {
        return m.rbegin()->first;
    }
    
    int popMax() {
        auto r = m.rbegin()->first;
        l.erase(m[r].back());
        m[r].pop_back();
        if (m[r].empty())
            m.erase(r);
        return r;
    }
    
    // Doubly LinkedList based solution.
    list<int> l;
    typedef  list<int>::iterator addr_t;
    map<int, vector<addr_t>> m;
};