Leetcode MaxStack in Python

Question

My solution to Leetcode MaxStack in python in Python. I have two solution one is using linked-list and another one is using list.

Design a max stack that supports push, pop, top, peekMax and popMax.

push(x) -- Push element x onto stack.
pop() -- Remove the element on top of the stack and return it.
top() -- Get the element on the top.
peekMax() -- Retrieve the maximum element in the stack.
popMax() -- Retrieve the maximum element in the stack, and remove it. If you find more than one maximum elements, only remove the top-most one.

Using list:

class MaxStack_list:
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.head = []
        self.max_val = None

    def push(self, x):
        """
        Push element x onto stack.
        :type x: int
        :rtype: void
        """
        self.head.append(x)
        self.max_val = max(max(self.head), x)

    def pop(self):
        """
        Removes the element on top of the stack and returns that element.
        :rtype: int
        """
        self.head.pop()
        self.max_val = max(self.head)

    def top(self):
        """
        Get the top element.
        :rtype: int
        """
        return self.head[-1] if self.head else None

    def peekMax(self):
        """
        Retrieve the maximum element in the stack.
        :rtype: int
        """
        return self.max_val

    def popMax(self):
        """
        Retrieve the maximum element in the stack, and remove it. If you find more than one maximum elements, only remove the top-most one.
        :rtype: void
        """
        v = self.max_val
        self.head.remove(self.max_val)
        self.max_val = max(self.head)
        return v

Using LinkedList

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


class MaxStack:
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.head = None
        self.max_val = None

    def push(self, x):
        """
        Push element x onto stack.
        :type x: int
        :rtype: void
        """
        if self.head:
            n = Node(x)
            n.next = self.head
            self.head = n
        else:
            self.head = Node(x)
        self.max_val = max(x, self.max_val) if self.max_val or self.max_val == 0 else x

    def pop(self):
        """
        Removes the element on top of the stack and returns that element.
        :rtype: int
        """
        rtn = None
        if self.head:
            rtn = self.head.val
            self.head = self.head.next

        head = self.head
        v = head.val if head else None
        while head:
            v = max(v, head.val)
            head = head.next
        self.max_val = v
        return rtn

    def top(self):
        """
        Get the top element.
        :rtype: int
        """
        if self.head:
            return self.head.val

    def peekMax(self):
        """
        Retrieve the maximum element in the stack.
        :rtype: int
        """
        return self.max_val

    def popMax(self):
        """
        Retrieve the maximum element in the stack, and remove it. If you find more than one maximum elements, only remove the top-most one.
        :rtype: void
        """
        prev, cur = None, self.head
        while cur:
            if cur.val == self.max_val and cur == self.head:
                self.head = cur.next
                break
            elif cur.val == self.max_val:
                prev.next = cur.next
                break

            prev, cur = cur, cur.next

        cur = self.head
        tmp = self.max_val
        v = cur.val if cur else None
        while cur:
            if cur:
                v = max(v, cur.val)
            cur = cur.next
        self.max_val = v
        return tmp

Answer 1

The return type for popMax should be int not void.

When you push a new item, your max can only increase above the current maximum. So instead of

self.max_val = max(max(self.head), x)

which is \$O(n)\$, you should have

self.max_val = max(self.max_val, x)

which is \$O(1)\$.

Theif self.max_val or self.max_val == 0 condition is better written as if self.max_val is not None.