本文旨在指导读者如何使用单调栈这一数据结构，将一个时间复杂度为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]`。 ### 原始代码及其时间复杂度分析 提供的原始代码使用队列 `queue.Queue()` 来实现编码逻辑，其核心思想是遍历队列中的每个元素，并在队列的剩余部分中查找第一个更大的元素。 “`python import queue q = queue.Queue() a = [4, 3, 7, 3, 2, 8, 6, 1, 10, 3] for i in a: q.put(i) encoded = [] while q: current = q.get() for i in range(q.qsize()): if current < q.queue[i]: encoded.append(q.queue[i] + current) break print(encoded)

这段代码的时间复杂度是 o(n²)，因为对于队列中的每个元素，都需要遍历队列的剩余部分来寻找更大的元素。

使用单调栈优化

单调栈是一种特殊的栈结构，其内部元素保持单调递增或单调递减的顺序。利用单调栈，我们可以高效地找到数组中每个元素后面第一个更大的元素。

算法思路：

1. 创建一个空栈 s 用于存储数组元素的索引。
2. 遍历数组 a，对于每个元素 x，执行以下操作：
   • 如果栈不为空，并且当前元素 x 大于栈顶元素对应的值 a[s[-1]]，则循环执行以下操作：
     • 将栈顶元素弹出，并将其对应编码后的值更新为栈顶元素的值加上当前元素 x。
   • 将当前元素的索引 i 压入栈中。
3. 遍历结束后，栈中剩余的元素表示它们后面没有更大的元素，因此它们的编码值保持不变（即加上自身，但由于初始化时已经复制了数组，所以无需额外处理）。

Python 代码实现：

a = [4, 3, 7, 3, 2, 8, 6, 1, 10, 3]  encoded = a[:]  # 创建一个a的副本，用于存放编码后的值 s = []  # 创建一个空栈 for i, x in enumerate(a):     while s and x > a[s[-1]]:         encoded[s.pop()] += x     s.append(i)  print(encoded)

代码解释：

立即学习“Python免费学习笔记（深入）”；

快转字幕

新一代 AI 字幕工作站，为创作者提供字幕制作、学习资源、会议记录、字幕制作等场景，一键为您的视频生成精准的字幕。

362

查看详情

• encoded = a[:] 创建了数组 a 的一个副本，这是为了避免直接修改原始数组。
• s = [] 初始化一个空栈，用于存储数组元素的索引。
• for i, x in enumerate(a): 遍历数组 a，同时获取元素的索引 i 和值 x。
• while s and x > a[s[-1]]: 这是一个循环，用于处理栈中元素。如果栈不为空，并且当前元素 x 大于栈顶元素对应的值 a[s[-1]]，则说明找到了栈顶元素后面第一个更大的元素。
• encoded[s.pop()] += x 将栈顶元素弹出，并将其对应编码后的值更新为栈顶元素的值加上当前元素 x。
• s.append(i) 将当前元素的索引 i 压入栈中。

时间复杂度分析：

虽然代码中有一个 while 循环，但每个元素最多入栈一次，也最多出栈一次。因此，内层 while 循环的总执行次数不会超过 n，其中 n 是数组的长度。所以，整个算法的时间复杂度是 O(n)。

注意事项和总结

• 单调栈是一种非常有用的数据结构，可以用于解决很多与寻找“下一个更大/更小元素”相关的问题。
• 在使用单调栈时，需要仔细考虑栈中应该存储元素的值还是索引，以及如何维护栈的单调性。
• 上述代码中，我们存储的是元素的索引，这样可以方便地更新 encoded 数组。
• 理解单调栈的关键在于理解其单调性如何帮助我们找到目标元素，并避免不必要的比较。

通过使用单调栈，我们将原始代码的时间复杂度从 O(n²) 降低到 O(n)，显著提高了代码的效率。这个案例展示了如何利用合适的数据结构和算法来优化代码的性能。