【Lintcode】859. Max Stack

题目地址：

https://www.lintcode.com/problem/max-stack/description

要求实现一个栈，可以实现下面操作：
1、push(x) – 将元素$x$添加到栈中。
2、pop() – 删除栈中最顶端的元素并将其返回。
3、top() – 返回栈中最顶端的元素。
4、peekMax() – 返回栈中最大的元素。
5、popMax() – 返回栈中最大的元素，并将其删除。如果有多于一个最大的元素，只删除最靠近顶端的一个元素。

可以用TreeMap + 双向链表来实现，双向链表负责实现栈的后进先出的逻辑，TreeMap来实现找最大值的逻辑。TreeMap的key里存栈里的数值，value存key对应的链表节点（由于栈里可能有重复数值，所以value应该存一个链表节点的list，并且应该将链表节点按照进栈的顺序加入这个list）。执行push、pop和top的时候只需要维护两个数据结构的定义即可，对于关于max的操作，只需从TreeMap里找最大值即可，复杂度是$O(1)$的。代码如下：

import java.util.ArrayList;
import java.util.List;
import java.util.TreeMap;

class MaxStack {
    
    class Node {
        int val;
        Node prev, next;
    
        public Node(int val) {
            this.val = val;
        }
    }
    
    private Node head, tail;
    private TreeMap<Integer, List<Node>> treeMap;
    
    public MaxStack() {
        // do intialization if necessary
        head = new Node(-1);
        tail = new Node(-1);
        head.next = tail;
        tail.prev = head;
        treeMap = new TreeMap<>();
    }
    
    /*
     * @param number: An integer
     * @return: nothing
     */
    public void push(int x) {
        // write your code here
        Node node = tail.prev;
        node.next = new Node(x);
        node.next.prev = node;
        node.next.next = tail;
        tail.prev = node.next;
        
        treeMap.putIfAbsent(x, new ArrayList<>());
        treeMap.get(x).add(tail.prev);
    }
    
    public int pop() {
        // write your code here
        tail = tail.prev;
        tail.next = null;
        int val = tail.val;
    
        List<Node> list = treeMap.get(val);
        list.remove(list.size() - 1);
        // 注意如果val在栈中已经不存在了，需要在treeMap里也将其删去
        if (list.isEmpty()) {
            treeMap.remove(val);
        }
        
        return val;
    }
    
    /*
     * @return: An integer
     */
    public int top() {
        // write your code here
        return tail.prev.val;
    }
    
    /*
     * @return: An integer
     */
    public int peekMax() {
        // write your code here
        return treeMap.lastKey();
    }
    
    /*
     * @return: An integer
     */
    public int popMax() {
        // write your code here
        int x = treeMap.lastKey();
        List<Node> list = treeMap.get(x);
        Node node = list.remove(list.size() - 1);
        node.prev.next = node.next;
        node.next.prev = node.prev;
        if (list.isEmpty()) {
            treeMap.remove(x);
        }
        
        return x;
    }
}

top()和peekMax()时间复杂度$O(1)$，其余时间复杂度都是$O(\log n)$，$n$为当前栈里元素个数，空间$O(n)$。