CustomStack.java

private static class CustomStackImpl<T extends Comparable<? extends T>> {
        private final Node<T> EMPTY = new Node<>(null, null, null, -1);

        private int counter = 0;
        private Node<T> head = EMPTY;
        private final TreeMap<T, SortedSet<Node<T>>> orderedMap = new TreeMap<>();

        public T push(T x) {
            int index = counter++;
            head = head.add(x, index);
            orderedMap.computeIfAbsent(x, key -> new TreeSet<>()).add(head);
            return x;
        }

        public T pop() {
            Node<T> removed = head;
            this.head = head.remove();
            orderedMap.get(removed.value).remove(removed);
            return removed.value;
        }

        public T top() {
            return head.value;
        }

        public T pickMax() {
            return orderedMap.firstEntry().getKey();
        }

        public T popMax() {
            Map.Entry<T, SortedSet<Node<T>>> entry = orderedMap.firstEntry();
            Node<T> removed;
            if (entry.getValue().size() == 1) {
                removed = entry.getValue().first();
                orderedMap.remove(entry.getKey());
            } else {
                removed = entry.getValue().last();
                entry.getValue().remove(removed);
            }

            if (removed == this.head)
                this.head = removed.remove();

            return removed.value;
        }

        private final class Node<K> implements Comparable<Node<? extends K>> {
            private Node<K> prev;
            private Node<K> next;
            private K value;
            private final int index;

            public Node(Node<K> prev, Node<K> next, K value, int index) {
                this.value = value;
                this.prev = prev;
                this.next = next;
                this.index = index;
            }

            public Node<K> add(K x, int index) {
                next = new Node<>(this, EMPTY, x, index);
                return next;
            }

            public Node<K> remove() {
                this.prev.next = this.next;
                this.next.prev = this.prev;
                if (this.prev.next == EMPTY)
                    return this.prev;
                else
                    return this.prev.next;
            }


            @Override
            public int compareTo(Node<? extends K> o) {
                return Integer.compare(this.index, o.index);
            }
        }
    }