S m j i/stack/min const time

data_structures/stacks/min_constant_time.py

@@ -0,0 +1,72 @@
+"""
+Given an set of numbers in a stack,
+find the minimum value from the stack at O(1)
+
+Problem: https://leetcode.com/problems/min-stack/description/
+"""
+
+stack = []
+min_stack = []
+
+
+def push(value: int):
+    """
+    Push into the main stack and track the minimum.
+    If the value to insert < minimum, then push to min stack
+    Returns None
+    """
+    if len(stack) == 0:
+        min_stack.append(value)
+        stack.append(value)
+        return
+
+    if value < min_stack[-1]:
+        min_stack.append(value)
+    stack.append(value)
+
+
+def pop():
+    """
+    Pop from the stack.
+    If the popped value is the same as the min stack top,
+    pop from the min stack as well
+
+    Returns None
+    """
+    if len(stack) == 0:
+        print("Nothing on stack")
+        return
+
+    top = stack.pop()
+    if len(min_stack) > 0 and top == min_stack[-1]:
+        min_stack.pop()
+
+
+def get_min():
+    """
+    Return the minimum element of the main stack by
+    returning the top of the minimum stack
+
+    Returns the minium element (int)
+
+    >>> push(10)
+    >>> push(20)
+    >>> push(5)
+    >>> push(30)
+    >>> push(1)
+    >>> get_min()
+    1
+    >>> pop()
+    >>> get_min()
+    5
+    >>> pop()
+    >>> get_min()
+    10
+    """
+    return min_stack.pop()
+
+
+if __name__ == "__main__":
+    from doctest import testmod
+
+    testmod()