Problem statement

https://binarysearch.com/problems/Maximum-Stack/

Solution

Equal to Leetcode 0716 Max Stack, but here we need fast operations.

Complexity

Time complexity for all operations except peek is O(log n), for peek it is O(1). Space is O(n).

Code

from sortedcontainers import SortedDict

class ListNode:
    def __init__(self, val=0, nxt=None, prv=None):
        self.val = val
        self.next = nxt
        self.prev = prv

class DoubleLinkedList:
    def __init__(self):
        self.head = ListNode(None)
        self.tail = ListNode(None)
        self.head.next = self.tail
        self.tail.prev = self.head

    def add(self, val):
        x = ListNode(val)
        x.next = self.tail
        x.prev = self.tail.prev
        self.tail.prev.next = x
        self.tail.prev = x
        return x

    def delete(self, node):
        node.prev.next = node.next
        node.next.prev = node.prev

class MaximumStack:
    def __init__(self):
        self.map = SortedDict()
        self.dll = DoubleLinkedList()

    def append(self, x):
        node = self.dll.add(x)
        if x not in self.map:
            self.map[x] = [node]
        else:
            self.map[x].append(node)

    def pop(self):
        val = self.dll.tail.prev.val
        self.dll.delete(self.dll.tail.prev)
        self.map[val].pop()
        if not self.map[val]: self.map.pop(val)
        return val

    def peek(self):
        return self.dll.tail.prev.val

    def max(self):
        return self.map.peekitem()[0]

    def popmax(self):
        max = self.max()
        node = self.map[max][-1]
        self.dll.delete(node)
        self.map[max].pop()
        if not self.map[max]: self.map.pop(max)
        return max

Solution 2

Another way is to use heaps with lazy deletions here.

  1. self.stk is stack of pairs (element, time)
  2. self.pq is a priority queue of pairs (-element, -time)
  3. self.time_to_val is correspondence between times and values.
  4. fix is function to remove irrelevant elements form stack and heap, we can call it lazy updates.

Complexity

It is O(log n) for all operations except peek and max

Code

### borrowed code
class MaximumStack:
    def __init__(self):
        self.time = 0
        self.time_to_val = {}
        self.stk = []
        self.pq = []

    def append(self, val):
        self.time += 1
        self.time_to_val[self.time] = val
        self.stk.append((val, self.time))
        heappush(self.pq, (-val, -self.time))

    def fix(self, invalid_time):
        del self.time_to_val[invalid_time]
        while self.stk and (self.stk[-1][1] not in self.time_to_val):
            self.stk.pop()
        while self.pq and (-self.pq[0][1] not in self.time_to_val):
            heappop(self.pq)

    def peek(self):
        return self.stk[-1][0]

    def max(self):
        return -self.pq[0][0]

    def pop(self):
        entry = self.stk.pop()
        self.fix(entry[1])
        return entry[0]

    def popmax(self):
        entry = heappop(self.pq)
        self.fix(-entry[1])
        return -entry[0]