如何使用单调栈优化 Python 代码的时间复杂度

本文旨在指导读者如何使用单调栈这一数据结构，将原本时间复杂度为 O(n²) 的 Python 代码优化至 O(n)。通过具体示例和详细解释，我们将展示如何利用单调栈高效地找到数组中每个元素的下一个更大元素，从而提升算法性能。 ### 问题描述 给定一个数组，目标是将每个元素替换为该元素与数组中其后第一个更大元素的和。如果某个元素之后没有更大的元素，则该元素的值保持不变。例如，对于输入数组 `[4, 3, 7, 3, 2, 8, 6, 1, 10, 3]`，期望的输出是 `[11, 10, 15, 11, 10, 18, 16, 11, 10, 3]`。 ### 原始代码及复杂度分析 提供的原始代码使用了嵌套循环，导致时间复杂度为 O(n²)。外层循环遍历数组中的每个元素，内层循环则查找该元素之后的第一个更大元素。这种方法效率较低，尤其是在处理大型数组时。 ### 优化方案：单调栈 单调栈是一种特殊的栈结构，其内部元素保持单调递增或单调递减的顺序。利用单调栈，我们可以在 O(n) 的时间复杂度内找到数组中每个元素的下一个更大元素。 #### 单调栈的工作原理 1. **初始化：** 创建一个空栈 `s`，用于存储数组元素的索引。 2. **遍历数组：** 从头到尾遍历数组 `a`。 3. **维护单调性：** 对于当前元素 `x`，执行以下操作： * 如果栈 `s` 不为空，并且 `x` 大于 `a[s[-1]]`（栈顶元素对应的值），则循环执行以下操作： * 弹出栈顶元素 `index`。 * 将 `encoded[index]` 更新为 `encoded[index] + x`。 * 将当前元素的索引 `i` 压入栈 `s`。 4. **处理剩余元素：** 遍历结束后，栈 `s` 中可能还存在一些元素，这些元素在数组中没有找到更大的元素，因此它们的值保持不变。 #### 代码实现 ```python def encode_array(a): """ 使用单调栈优化数组编码过程。 Args: a: 输入数组。 Returns: 编码后的数组。 """ encoded = a[:] # 创建数组的副本，避免修改原始数组 s = [] # 初始化单调栈 for i, x in enumerate(a): while s and x > a[s[-1]]: encoded[s.pop()] += x s.append(i) return encoded # 示例 a = [4, 3, 7, 3, 2, 8, 6, 1, 10, 3] encoded = encode_array(a) print(encoded) # 输出: [11, 10, 15, 11, 10, 18, 16, 11, 10, 3]

代码解释

• encoded = a[:] 创建了输入数组 a 的一个副本，这样修改 encoded 不会影响原始数组。
• s = [] 初始化了单调栈。
• enumerate(a) 用于同时获取数组的索引和值。
• while s and x > a[s[-1]] 循环确保栈的单调性。当遇到比栈顶元素更大的元素时，不断弹出栈顶元素，直到栈为空或者栈顶元素大于等于当前元素。
• encoded[s.pop()] += x 将栈顶元素对应的值更新为与当前元素的和。
• s.append(i) 将当前元素的索引压入栈中。

复杂度分析

• 时间复杂度： O(n)。虽然代码中包含一个 while 循环，但每个元素最多入栈一次，出栈一次，因此总的时间复杂度为 O(n)。
• 空间复杂度： O(n)。单调栈最多存储 n 个元素的索引。

注意事项

• 单调栈适用于解决寻找数组中下一个更大/更小元素的问题。
• 在实际应用中，可以根据具体需求选择单调递增栈或单调递减栈。
• 使用单调栈时，需要注意维护栈的单调性，确保算法的正确性。

总结

通过使用单调栈，我们可以将原本时间复杂度为 O(n²) 的代码优化至 O(n)，显著提升算法的性能。单调栈是一种强大的数据结构，在解决与数组元素大小关系相关的问题时非常有用。理解单调栈的工作原理和应用场景，可以帮助我们编写更高效的 Python 代码。