LeetCode 716. Max Stack solution: O(log N) using std::set (no linked lists)

leetcode 716 Max Stack.md

https://leetcode.com/problems/max-stack

https://www.lintcode.com/problem/859

Design a max stack that supports:

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.

---

std::set in C++ stores elements in sorted order, and allows to add/find/delete any element by it's value (which acts as key). It's similar to std::map, but no additional data (a.k.a value) is associated with the keys.

For this task, don't think about std::set in terms of uniqueness (while technically value are still unique). Think of it like just sorted container.

Side note: what is called "set" and "dict" in other languages, in C++ corresponds to "unordered_set" and "unordered_map".

So what I do: together with each number v I store "timestamp" t. The timestamp is just integer which increases by 1 on each push. This will allow to distinguish numbers with same value.

Then i store each pair <v, t> in two containers: one is sorted by v (and if two elements have the same v, then they're sorted by t), and the second is sorted by t. To do that, I just declare set of pairs: default comparator compares by first item and then by second item.

All the rest you can see in the code. rbegin is the last (biggest) element.

#include <set>
#include <iostream>

template<class V>
class MaxStackAny {
    using uint = unsigned int;
    uint last_t = 0;
    std::set<std::pair<uint, V>> by_t;
    std::set<std::pair<V, uint>> by_v;

public:
    void push(V v) {
        uint t = ++last_t;
        by_t.insert({t, v});
        by_v.insert({v, t});
    }

    V pop() {
        auto t = by_t.rbegin()->first;
        auto v = by_t.rbegin()->second;
        by_t.erase({t, v});
        by_v.erase({v, t});
        return v;
    }

    V top() {
        return by_t.rbegin()->second;
    }

    V peekMax() {
        return by_v.rbegin()->first;
    }

    V popMax() {
        auto v = by_v.rbegin()->first;
        auto t = by_v.rbegin()->second;
        by_t.erase({t, v});
        by_v.erase({v, t});
        return v;
    }
};

using MaxStack = MaxStackAny<int>;

int main() {
    MaxStackAny<std::string> s;
    s.push("20");
    s.push("10");
    std::cout << s.peekMax() << std::endl << std::endl;

    MaxStackAny<int> i;
    i.push(20);
    i.push(10);
    std::cout << i.peekMax() << std::endl << std::endl;

    MaxStack f;
    f.push(5);
    f.push(1);
    f.push(5);
    std::cout << f.top() << std::endl;
    std::cout << f.popMax() << std::endl;
    std::cout << f.top() << std::endl;
    std::cout << f.peekMax() << std::endl;
    std::cout << f.pop() << std::endl;
    std::cout << f.top() << std::endl;
}

Here push, pop/popMax are O(log N), and top/peekMax are O(1) according to std::set API and as there are no loops in my code.

The code can be further optimized:

• In pop we can erase the element from by_t without searching (like we currently do by passing the pair) but by iterator. Same for popMax and by_v
• As timestamps in by_t are unique, we can replace set with map<uint, V>
• For each v in by_v it would be enough to store min_t and max_t, so we can replace set with map<T, pair<uint, uint>>, but this will make implementation some harder But I think the current is clean (see symmetry of by_t and by_v, also pop and popMax) which I think is important, while optimizations not necessary will make it faster.