题面及题解

原题链接

给你一个整数数组 nums 和一个正整数 k 。请你统计有多少满足「 nums 中的最大元素」至少出现 k 次的子数组，并返回满足这一条件的子数组的数目。
子数组是数组中的一个连续元素序列。

题解

题目要求的是 nums 中的最大元素，这是固定的，可以对原数组使用滑动窗口解决，难度并不高。

def countSubarrays(nums: list[int], k: int) -> int:
    mx = max(nums)
    left = right = cntMx = ans = 0
    for right in range(len(nums)):
        if nums[right] == mx:
            cntMx += 1
        while cntMx == k:
            if nums[left] == mx:
                cntMx -= 1
            left += 1
        ans += left
    return ans

扩展

原题并不困难，但是如果读错了题，将「nums 中的最大元素」理解为了「子数组的中最大元素」，那么就不能简单地用滑动窗口来解决了。

事实上，评论区表明有很多人这样子理解了

那么如果要求的是「子数组的中最大元素」，应该如何解决呢？

错误思路

一开始我联想到了力扣中滑动窗口最大值这一题。该题使用单调栈来维护窗口内的元素大小排序。

那么放到这里来的话，可以在窗口的基础上也用一个单调栈来维护窗口内的元素，同时用一个字典维护栈内元素的数量。

在窗口右指针滑动时将元素压入栈，并更新字典中元素数量。当「栈底元素（即窗口内的最大值）数量大于等于k」时，移动窗口左指针。将左指针移出的元素更新栈和字典，直至不满足条件。此时便可计算一轮满足条件的子数组的数量。

然而以上思路存在错误。

例如，对于此数组，[2, 1, 0, 2]：

窗口在移动时，2 将始终是窗口内的最大值，无法统计到最大值为 1 或者 0 的子数组情况。

究其原因，在这种情况下，我们维护的滑动窗口内的状态并不满足单调性，滑动过程中不能简单地根据左右指针计算答案数量。这是滑动窗口本身的适用范围。

解题思路

根据以上结论，我们应该思考如何满足滑动窗口内的单调性。

可以考虑将问题转换成类似原题的情形。即对于数组中的某个数 x，我们找出以 x 为最大值的一段连续区间 (l, r)。那么在这个范围内，我们可以统计出 x 的数量大于 k 的子数组。

当然，以上思路相当多元素仍然重复参与了计算，需要继续优化。

单调栈统计区间(l, r)

可以使用单调栈分别从左到右、从右到左统计每个元素的第一个大于该元素的位置 left、right。时间复杂度为 O(n)。

往右单向统计相同元素位置

可以知道，对于区域 (l, r) 内的每一个最大元素，其适用的区域是一致的。

那么我们遍历数组第一次遇到元素 x 时，根据left, right 快速得出其为最大值的区间后，可以在区间内向右找到所有值为 x 的位置，有了这些位置后就能计算最大值为 x 的满足条件的子数组的数量。此后在该区间内在遇到 x 就可以直接跳过了。

如何快速定位 x 的位置呢？我们可以修改单调栈的条件，统计每个元素第一个大于等于它的位置，就可以通过 right 数组快速跳转至相同元素或者区间右端点处。同时判断当前元素是否已处理，也只需要判断其left是否与其相等即可（有点类似并查集的思想）。

这样一来，遍历数组的过程中，每个元素最多访问2遍，时间复杂度为 O(n)。

实现

def numSubarrayBoundedMax(nums, k):
    n = len(nums)
    left = [-1] * n # left[i]表示i左边第一个大于等于nums[i]的位置
    right = [n] * n # right[i]表示i右边第一个大于等于nums[i]的位置
    stk = []
    for i in range(n):
        while stk and nums[stk[-1]] < nums[i]:
            stk.pop()
        if stk:
            left[i] = stk[-1]
        stk.append(i)
    stk = []
    for i in range(n - 1, -1, -1):
        while stk and nums[stk[-1]] < nums[i]:
            stk.pop()
        if stk:
            right[i] = stk[-1]
        stk.append(i)
    
    ans = 0
    for i, x in enumerate(nums):
        if left[i] != -1 and nums[left[i]] == x:
            '''
            如果nums[left[i]] == x，说明i的左边第一个大于等于nums[i]的数等于nums[i]
            那么i位置的元素在前面遍历到left[i]位置时已经处理，可以直接跳过
            '''
            continue
        keys = [left[i], i] # 存储区间边界以及内部值为x的位置
        ne_key = right[i]
        cnt = 1
        while ne_key != n and nums[ne_key] == x:
            '''
            通过while循环跳转至区间右边界，同时统计区间内x的位置
            '''
            cnt += 1
            keys.append(ne_key)
            ne_key = right[ne_key]
        keys.append(ne_key)
        if cnt >= k:
            for j in range(1, len(keys) - k):
                '''
                根据乘法原理从左往右统计x数量大于等于k的子数组的数目
                '''
                ans += (keys[j] - keys[j - 1]) * (keys[-1] - keys[j + k - 1])

    return ans