20 包含min函数的栈

题目描述

定义栈的数据结构，请在该类型中实现一个能够得到栈中所含最小元素的min函数（时间复杂度应为O(1)）。

解题思路

核心思想：使用两个栈，一个正常存数据，另一个存当前的最小值。

关键步骤：

1. push：数据栈正常push，最小栈push当前元素与栈顶的较小值
2. pop：两个栈同时pop
3. min：直接返回最小栈的栈顶

时间复杂度：所有操作都是O(1) 空间复杂度：O(n)

Python 实现

class StackWithMin:
    def __init__(self):
        self.data_stack = []  # 数据栈
        self.min_stack = []   # 最小栈

    def push(self, val):
        """
        入栈操作
        """
        self.data_stack.append(val)

        # 最小栈存当前最小值
        if not self.min_stack or val < self.min_stack[-1]:
            self.min_stack.append(val)
        else:
            self.min_stack.append(self.min_stack[-1])

    def pop(self):
        """
        出栈操作
        """
        if not self.data_stack:
            raise Exception("Stack is empty")

        self.min_stack.pop()
        return self.data_stack.pop()

    def min(self):
        """
        获取栈中最小元素
        """
        if not self.min_stack:
            raise Exception("Stack is empty")

        return self.min_stack[-1]

    def peek(self):
        """
        查看栈顶元素
        """
        if not self.data_stack:
            raise Exception("Stack is empty")

        return self.data_stack[-1]

    def is_empty(self):
        """判断栈是否为空"""
        return not self.data_stack


# 测试代码
if __name__ == "__main__":
    stack = StackWithMin()

    stack.push(3)
    print(f"push(3), min={stack.min()}")  # 3

    stack.push(4)
    print(f"push(4), min={stack.min()}")  # 3

    stack.push(2)
    print(f"push(2), min={stack.min()}")  # 2

    stack.push(1)
    print(f"push(1), min={stack.min()}")  # 1

    print(f"pop()={stack.pop()}, min={stack.min()}")  # 1, 2
    print(f"pop()={stack.pop()}, min={stack.min()}")  # 2, 3
    print(f"pop()={stack.pop()}, min={stack.min()}")  # 4, 3
    print(f"pop()={stack.pop()}")  # 3

# 输出:
# push(3), min=3
# push(4), min=3
# push(2), min=2
# push(1), min=1
# pop()=1, min=2
# pop()=2, min=3
# pop()=4, min=3
# pop()=3

图解

操作过程:

push(3):
data_stack: [3]
min_stack:  [3]  (第一个元素直接入)

push(4):
data_stack: [3, 4]
min_stack:  [3, 3]  (4 > 3, 存3)

push(2):
data_stack: [3, 4, 2]
min_stack:  [3, 3, 2]  (2 < 3, 存2)

push(1):
data_stack: [3, 4, 2, 1]
min_stack:  [3, 3, 2, 1]  (1 < 2, 存1)

pop():
data_stack: [3, 4, 2]
min_stack:  [3, 3, 2]  (两个栈同时出)
min() = 2

pop():
data_stack: [3, 4]
min_stack:  [3, 3]
min() = 3

pop():
data_stack: [3]
min_stack:  [3]
min() = 3

pop():
data_stack: []
min_stack:  []