在 O(1) 时间和 O(1) 额外空间中查找堆栈中的最大值

在问题简述中，我们需要实现一个数据结构，它满足以下三个功能：

1. 使用 push 将元素添加到堆栈中。
2. 使用 pop 从堆栈中弹出元素。
3. 使用 getMax 返回堆栈中的最大值。

解法

常规解法

一个朴素的想法是在堆栈中存储一个最大值，每次插入一个元素时与当前最大值进行比较，更新最大值。但是这样会有一个问题：在弹出最大值后，下一个最大值是什么？堆栈中可能还有比下一个最大值更大的元素，但在此处该元素已经被忽略了。

那么，我们是否可以在堆栈中存储每个元素的最大值呢？这个想法看起来比前一个更有趣。

堆栈是被 First In Last Out（先进后出）约束的一种数据结构，我们可以使用一个辅助堆栈来存储每个当前元素在主堆栈中对应的最大值。具体地，当在主堆栈中插入元素时，我们比较当前元素和辅助堆栈的顶部元素，将较大者压入辅助堆栈。在主堆栈顶部弹出元素时，我们也将辅助堆栈顶部的元素弹出。插入、删除和查询操作的时间复杂度均为 O(1)。

下面是示例代码：

class MaxStack:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.stack = []
        self.max_stack = []

    def push(self, x: int) -> None:
        self.stack.append(x)
        if not self.max_stack or self.max_stack[-1] <= x:
            self.max_stack.append(x)

    def pop(self) -> int:
        if not self.stack:
            return None
        if self.stack[-1] == self.max_stack[-1]:
            self.max_stack.pop()
        return self.stack.pop()

    def top(self) -> int:
        if not self.stack:
            return None
        return self.stack[-1]

    def getMax(self) -> int:
        if not self.max_stack:
            return None
        return self.max_stack[-1]

优化解法

我们可以发现，在每次插入元素时，我们仅在当前元素大于最大堆栈的顶部元素时才向最大堆栈中加入元素。这给了我们一些提示：我们可以使用一个变量来存储当前已知的最大元素，而非维护一个最大堆栈。

也就是说，我们可以在主堆栈中存储一个差值，差值等于当前元素与当前最大值的差，而非元素本身。在插入元素时，我们找到该元素与当前最大值之间的差值，将该差值压入堆栈中。若该差值是正数（即当前元素比当前最大值大），则更新最大值。

删除操作的处理方式相同。在堆栈中取出元素时，我们取出的不是真正的元素。而是最大值和当前元素之间的差值，更新当前最大值。查询操作直接返回当前最大值即可。

下面是示例代码：

class MaxStack:

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

    def push(self, x: int) -> None:
        if not self.stack:
            self.max_val = x
            self.stack.append(0)
        else:
            diff = x - self.max_val
            self.stack.append(diff)
            self.max_val = x if diff > 0 else self.max_val

    def pop(self) -> int:
        if not self.stack:
            return None
        diff = self.stack.pop()
        if diff > 0:
            res = self.max_val
            self.max_val = self.max_val + diff if self.stack else None
        else:
            res = self.max_val + diff
        return res

    def top(self) -> int:
        if not self.stack:
            return None
        diff = self.stack[-1]
        if diff > 0:
            return self.max_val
        else:
            return self.max_val + diff

    def getMax(self) -> int:
        return self.max_val if self.max_val is not None else None

总结

在本文中，我们通过两种不同的解决方案来实现在 O(1) 时间和 O(1) 额外空间中查找堆栈中的最大值。常规方案使用一个辅助堆栈来维护最大值，而优化的方案则使用一个变量来实现该功能。需要注意的是，两种解决方案在查询操作中的时间复杂度均为 O(1)。