列表元素与前序连续元素比较计数实现及O(n)复杂度优化咨询

原代码错误分析 #

你写的代码存在3个核心问题：

• count变量没有在每次遍历i的时候重置，会把前序i的计数结果累加进来，导致结果错误
• 内层用了while循环判断a[i] >= a[j]，满足条件时只会无限给count加1，不会跳出循环，属于逻辑误用，这里应该用if判断
• append操作放在了j的循环内，每个i会多次向li追加值，不符合每个i对应一个结果的需求

修正后的暴力实现代码 #

a = [1,2,3,2,4,6,1,2]
li = []
for i in range(len(a)):
    count = 0
    # 从i往前遍历，包括i自身
    for j in range(i, -1, -1):
        if a[i] >= a[j]:
            count += 1
        else:
            break
    li.append(count)
print(' '.join(map(str, li)))
# 输出：1 2 3 1 5 6 1 2

O(n)时间复杂度实现方案 #

这类求每个元素左侧最近的比它大的元素位置的问题，可以用单调栈算法优化到O(n)时间复杂度，每个元素只会入栈、出栈各一次：

a = [1,2,3,2,4,6,1,2]
li = []
# 栈内保存(元素值, 连续计数长度)，维持单调递增的特性
stack = []
for num in a:
    cnt = 1
    # 栈顶元素小于等于当前数，就弹出累加计数
    while stack and stack[-1][0] <= num:
        cnt += stack.pop()[1]
    stack.append((num, cnt))
    li.append(cnt)
print(' '.join(map(str, li)))
# 输出：1 2 3 1 5 6 1 2

嵌套循环优化说明 #

所有需要向前/向后查找第一个满足大小关系元素的场景，都可以用单调栈替代嵌套循环，把时间复杂度从O(n²)降到O(n)，常见适用场景包括接雨水、柱状图中最大矩形、下一个更大元素等经典算法题。

内容的提问来源于stack exchange，提问作者adarsh augustine