diff --git a/data_structures/stacks/min_constant_time.py b/data_structures/stacks/min_constant_time.py new file mode 100644 index 000000000000..91462e0452fa --- /dev/null +++ b/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()