如何优化Python代码时间复杂度？求O(n log 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)（比你预期的O(n log n)更高效），核心思路是维护一个单调递减的栈，从数组末尾向前遍历：

1. 初始化一个空栈，用于存储后续元素中可能成为“下一个更大值”的候选
2. 从数组最后一个元素开始向前遍历：
   • 若栈不为空且栈顶元素 ≤ 当前元素，弹出栈顶（因为当前元素比它大，它不可能成为前面元素的下一个更大值）
   • 此时，栈顶元素就是当前元素的下一个更大值（栈为空则没有更大值）
   • 计算编码值：栈为空则取当前元素本身，否则取当前元素 + 栈顶元素
   • 将当前元素压入栈，维护栈的单调递减性

实现代码 #

def encode_array(arr):
    stack = []
    encoded = []
    # 从后往前遍历数组
    for num in reversed(arr):
        # 弹出栈中所有小于等于当前数的元素
        while stack and stack[-1] <= num:
            stack.pop()
        # 计算编码值
        if stack:
            encoded.append(num + stack[-1])
        else:
            encoded.append(num)
        # 当前数入栈
        stack.append(num)
    # 因为是从后往前遍历，需要反转得到正确顺序
    return encoded[::-1]

# 测试示例
a = [4,3,7,3,2,8,6,1,10,3]
print(encode_array(a))  # 输出: [11, 10, 15, 11, 10, 18, 16, 11, 10, 3]

复杂度说明 #

每个元素最多入栈和出栈各一次，总操作次数为O(n)，空间复杂度为O(n)（最坏情况栈存储所有元素）。

内容的提问来源于stack exchange，提问作者Mike F