[leetcode #155] Min Stack

Problem

Design a stack that supports push, pop, top, and retrieving the minimum element in constant time.

Implement the MinStack class:

・ MinStack() initializes the stack object.
・ void push(int val) pushes the element val onto the stack.
・ void pop() removes the element on the top of the stack.
・ int top() gets the top element of the stack.
・ int getMin() retrieves the minimum element in the stack.

Example 1:

Input
["MinStack","push","push","push","getMin","pop","top","getMin"]
[[],[-2],[0],[-3],[],[],[],[]]

Output
[null,null,null,null,-3,null,0,-2]

Explanation
MinStack minStack = new MinStack();
minStack.push(-2);
minStack.push(0);
minStack.push(-3);
minStack.getMin(); // return -3
minStack.pop();
minStack.top();    // return 0
minStack.getMin(); // return -2

Constraints:

・ -2³¹ <= val <= 2³¹ - 1
・ Methods pop, top and getMin operations will always be called on non-empty stacks.
・ At most 3 * 10⁴ calls will be made to push, pop, top, and getMin.

Idea

간단해보이지만 약간 생각을 해야 하는 문제다. MinStack에 최소값을 따로 저장해뒀다가 push가 될 때만 비교해서 리턴하면 틀릴 수밖에 없다. pop 연산시 최소값이 빠져나가면 getMin에서 틀린 값을 리턴하기 때문이다.

이를 해결하기 위해 두 가지 자료구조를 이용했다. 기본 구조로 Stack을 사용했고, 최소값을 리턴할 용도로 TreeMap을 사용했다.

Stack에는 MinStack의 기본 연산을 적용하고, TreeMap에는 stack에 들어간 값을 key로, 해당 값의 빈도 수를 value로 넣었다. getMin에서는 TreeMap이 지원하는 ceilingKey를 호출하면 최소값을 간단하게 구할 수 있다.

제약 조건 중 최소값을 O(constant)로 리턴해야 하는 게 있는데, treeMap의 ceilingKey가 O(constant)가 맞는 걸까? 결과를 보면 아닌 거 같은데, 다른 사람들의 풀이도 참고해야겠다.

Solution

class MinStack {
    Stack<Integer> stack;
    TreeMap<Integer, Integer> treeMap;

    public MinStack() {
        stack = new Stack<>();
        treeMap = new TreeMap<Integer, Integer>();
    }

    public void push(int val) {
        stack.push(val);
        if (!treeMap.containsKey(val))
            treeMap.put(val, 0);
        treeMap.put(val, 1 + treeMap.get(val));
    }

    public void pop() {
        int val = stack.pop();
        int cnt = treeMap.get(val);
        if (cnt == 1){
            treeMap.remove(val);
        } else {
            treeMap.put(val, cnt - 1);
        }
    }

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

    public int getMin() {
        return treeMap.ceilingKey(Integer.MIN_VALUE);
    }
}

Reference

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