LeetCode 716 - Max Stack

The problem asks us to design a stack data structure, MaxStack, which behaves like a normal stack with additional operations for efficiently accessing and removing the maximum element.

leetcodehardlinked-liststackdesigndoubly-linked-listordered-set

LeetCode Problem 716

Difficulty: 🔴 Hard
Topics: Linked List, Stack, Design, Doubly-Linked List, Ordered Set

Solution

Problem Understanding

The problem asks us to design a stack data structure, MaxStack, which behaves like a normal stack with additional operations for efficiently accessing and removing the maximum element. Specifically, in addition to the standard stack operations push, pop, and top, we need to implement peekMax, which returns the maximum element without removing it, and popMax, which removes the top-most occurrence of the maximum element. The inputs are sequences of integers pushed onto the stack, and the outputs are the results of the various stack operations.

Constraints indicate that each integer can be between -10^7 and 10^7 and there can be up to 10^5 operations. The guarantees ensure that pop, top, peekMax, and popMax will never be called on an empty stack, so we do not need to handle empty-stack exceptions. Important edge cases include having multiple occurrences of the maximum element, negative values, and sequences where popMax is called repeatedly.

The challenge is to implement these operations efficiently. While top must be O(1), push, pop, peekMax, and popMax should be as close to O(log n) as possible. Naive solutions may not meet these requirements.

Approaches

A brute-force approach is to use a standard stack and compute the maximum by scanning the stack each time peekMax or popMax is called. This guarantees correctness but is too slow: peekMax and popMax would become O(n), which is not acceptable for 10^5 operations.

The optimal solution uses a combination of a doubly-linked list to represent the stack and a balanced tree or sorted list to track elements by value. This allows constant-time operations for top and logarithmic-time operations for push, pop, peekMax, and popMax. The key insight is maintaining references from each value in the sorted structure to its nodes in the linked list so we can remove them efficiently.

Approach Time Complexity Space Complexity Notes
Brute Force O(n) for peekMax/popMax, O(1) for push/pop/top O(n) Simple stack, linear scan for max
Optimal O(log n) for push/pop/peekMax/popMax, O(1) for top O(n) Uses doubly-linked list + sorted map/tree of nodes for fast access

Algorithm Walkthrough

  1. Use a doubly-linked list to maintain the stack in order. This allows O(1) push, pop, and top operations.
  2. Maintain a sorted map (Python: SortedDict from sortedcontainers) mapping each value to a list of linked list nodes containing that value. This lets us access the maximum value in O(log n) time.
  3. For push(x), create a new node at the tail of the linked list and add it to the sorted map under the key x.
  4. For pop(), remove the tail node from the linked list and remove the reference from the sorted map. Return the node value.
  5. For top(), simply return the value of the tail node.
  6. For peekMax(), query the sorted map for the largest key. This gives O(log n) access to the maximum value without removing it.
  7. For popMax(), find the largest key in the sorted map and remove the most recently added node for that value from the linked list and the map. If the list of nodes for that value becomes empty, remove the key from the map.

Why it works: The doubly-linked list preserves stack order and allows constant-time removal of arbitrary nodes. The sorted map keeps track of the current maximum efficiently. Combining these two ensures all operations meet the desired time complexity.

Python Solution

from sortedcontainers import SortedDict

class Node:
    def __init__(self, val: int):
        self.val = val
        self.prev = None
        self.next = None

class MaxStack:

    def __init__(self):
        self.head = Node(0)  # dummy head
        self.tail = Node(0)  # dummy tail
        self.head.next = self.tail
        self.tail.prev = self.head
        self.map = SortedDict()  # val -> list of nodes

    def push(self, x: int) -> None:
        node = Node(x)
        node.prev = self.tail.prev
        node.next = self.tail
        self.tail.prev.next = node
        self.tail.prev = node
        if x not in self.map:
            self.map[x] = []
        self.map[x].append(node)

    def pop(self) -> int:
        node = self.tail.prev
        self._remove(node)
        return node.val

    def top(self) -> int:
        return self.tail.prev.val

    def peekMax(self) -> int:
        return self.map.peekitem(-1)[0]

    def popMax(self) -> int:
        max_val, nodes = self.map.peekitem(-1)
        node = nodes.pop()
        if not nodes:
            del self.map[max_val]
        self._remove(node)
        return node.val

    def _remove(self, node: Node):
        node.prev.next = node.next
        node.next.prev = node.prev

The Python solution uses a doubly-linked list for stack order and a SortedDict to efficiently track maximum values. Each push appends to the tail, and _remove handles arbitrary node removal efficiently.

Go Solution

package main

import (
    "container/list"
    "sort"
)

type MaxStack struct {
    stack *list.List
    mapVal map[int][]*list.Element
    sortedKeys []int
}

func Constructor() MaxStack {
    return MaxStack{
        stack: list.New(),
        mapVal: make(map[int][]*list.Element),
        sortedKeys: []int{},
    }
}

func (this *MaxStack) Push(x int)  {
    node := this.stack.PushBack(x)
    this.mapVal[x] = append(this.mapVal[x], node)
    if len(this.mapVal[x]) == 1 {
        this.sortedKeys = append(this.sortedKeys, x)
        sort.Ints(this.sortedKeys)
    }
}

func (this *MaxStack) Pop() int {
    node := this.stack.Back()
    val := node.Value.(int)
    this.stack.Remove(node)
    nodes := this.mapVal[val]
    this.mapVal[val] = nodes[:len(nodes)-1]
    if len(this.mapVal[val]) == 0 {
        delete(this.mapVal, val)
        this.sortedKeys = this.sortedKeys[:len(this.sortedKeys)-1]
    }
    return val
}

func (this *MaxStack) Top() int {
    return this.stack.Back().Value.(int)
}

func (this *MaxStack) PeekMax() int {
    return this.sortedKeys[len(this.sortedKeys)-1]
}

func (this *MaxStack) PopMax() int {
    maxVal := this.PeekMax()
    nodes := this.mapVal[maxVal]
    node := nodes[len(nodes)-1]
    this.stack.Remove(node)
    this.mapVal[maxVal] = nodes[:len(nodes)-1]
    if len(this.mapVal[maxVal]) == 0 {
        delete(this.mapVal, maxVal)
        this.sortedKeys = this.sortedKeys[:len(this.sortedKeys)-1]
    }
    return maxVal
}

The Go implementation uses container/list for a doubly-linked list and a map from values to slices of list elements. Since Go lacks a built-in balanced tree, we maintain a sortedKeys slice. The logic mirrors Python but requires manual slice sorting.

Worked Examples

For the input sequence:

["push 5", "push 1", "push 5", "top", "popMax", "top", "peekMax", "pop", "top"]

Step by step:

Operation Stack Max Map Return
push 5 [5] {5:[node]} -
push 1 [5,1] {5:[node],1:[node]} -
push 5 [5,1,5] {5:[node,node],1:[node]} -
top [5,1,5] ... 5
popMax [5,1] {5:[node],1:[node]} 5
top [5,1] ... 1
peekMax [5,1] ... 5
pop [5] {5:[node]} 1
top [5] ... 5

Complexity Analysis

Measure Complexity Explanation
Time O(log n) push/pop/peekMax/popMax, O(1) top Sorted map ensures max access in O(log n), doubly-linked list allows O(1) removal
Space O(n) Storage in both linked list and sorted map for all nodes

The space complexity is linear in the number of elements because each element exists in both the stack and the mapping structure.

Test Cases

# test cases
obj = MaxStack()
obj.push(5)
obj.push(1)
obj.push(5)
assert obj.top() ==