Max Stack [LeetCode 716]

MaxStack.py

from sortedcontainers import SortedDict

class Element:
    def __init__(self, v):
        self.v = v      # Current value
        self.p = None   # Previous pointer
        self.n = None   # Next Pointer
        return

class MaxStack:

    def __init__(self):
        self.start = None
        self.end = None
        self.heap = SortedDict()
        return

    def push(self, x: int) -> None:
        
        # Make element
        element = Element(x)
        
        # Add to stack
        if self.start:
            self.end.n = element
            element.p = self.end
            self.end = element
        else:
            self.start = self.end = element
            
        # Add to heap
        if x in self.heap:
            self.heap[x].append(element)
        else:
            self.heap[x] = [element]
            
        return

    def pop(self) -> int:
        
        # Get the last element
        k = self.end.v
       
        # Remove from heap and stack
        self.removeStack(self.removeHeap(k))
        
        return k

    def top(self) -> int:
        return self.end.v # value of the top element

    def peekMax(self) -> int:
        return self.heap.keys()[-1] # value of the max element

    def popMax(self) -> int:
        
        # Get the max element
        k = self.heap.keys()[-1]
        
        # Remove from heap and stack
        self.removeStack(self.removeHeap(k))
            
        return k
    
    # A function to remove element w.r.t the key from heap
    def removeHeap(self, k: int):
        
        # Pop the element on top of the stack w.r.t key
        element = self.heap[k].pop()
        
        # Delete the key if all elements relate to key
        # is popped
        if not self.heap[element.v]:
            self.heap.pop(element.v, None)
        
        # return the popped element
        return element
    
    # A helper function to remove an element from stack
    def removeStack(self, curr: Element):
        
        # Get previous and next pointers of the
        # current element
        prev = curr.p
        next = curr.n
        
        # Case 1: only one single element in stack
        if not prev and not next:
            self.start = self.end = None
        
        # Case 2: removing at the start(bottom) of the stack
        elif not prev:
            next = self.start.n
            self.start = next
            if next:
                next.p = None
        
        # Case 3: removing at the end(top) of the stack
        elif not next:
            prev = self.end.p 
            self.end = prev
            if prev:
                prev.n = None
        
        # Case 4: removing the at middle of the stack
        else:
            next.p = curr.p
            prev.n = curr.n
        
        # Remove all pointers w.r.t curr element
        if curr:
            if curr.p:
                curr.p = None
            if curr.n:
                curr.n = None
            curr = None
            
        return


# Your MaxStack object will be instantiated and called as such:
# obj = MaxStack()
# obj.push(x)
# param_2 = obj.pop()
# param_3 = obj.top()
# param_4 = obj.peekMax()
# param_5 = obj.popMax()