计算最大元素出现至少 K 次的子数组数量滑动窗口

Question

问题

计算最大元素至少出现 K 次的子数组

给你一个整数数组 nums 和一个正整数 k。你需要返回在这个数组中，最大元素至少出现 k 次的子数组的数量。子数组是数组中连续的一段元素。

逻辑

• 遍历窗口，如果当前最大元素的出现次数少于 k，就继续扩大窗口。
• 如果当前最大元素的出现次数大于等于 k，就增加结果 ans，然后尝试缩小或滑动窗口。

代码

class Solution:

    def countSubarrays(self, nums: List[int], k: int) -> int:
        i = j = 0
        count = 0
        ans = 0 
        N = len(nums)
        max_element = max(nums)
    
        while j < N:
            # keep increasing the window
            if count < k:
                count+= nums[j] == max_element
                j+=1
            # slide/decrease the window
            else:
                ans+=1
                if nums[i] == max_element:
                    count-=1
                i+=1
    
        return ans
            

问题

可能遇到的问题包括：

• 我可能没有考虑到最后一个子数组，这个子数组可能需要缩短，或者至少应该被加入到 ans 变量中。
• 我没有正确滑动窗口，因为我可以向右也可以向左。比如说 nums = [1,3,2,3,3]，当前子数组是 [3,2,3]，我已经满足了 count >= k，所以我应该考虑 [3,2,3,3]，而不是只缩短当前子数组。
• 我知道可以用暴力方法来做一个循环（当 count>=k 时，增加 i 直到条件不再满足），但我觉得这并不能解决问题。这种情况也可能发生在一个条件判断中（如果 count >=k，就固定 j 直到条件不再满足），所以任何帮助都很感激。我希望能找到一个解决方案，当 count >=k 时，我们不增加 j，直到 count < k，然后再继续增加 j。
• 如果你需要一个具体的例子，这段代码在 nums=[1,3,2,3,3] 和 k = 2 时失败，输出是 2，但期望输出是 6。两个子数组是 [1, 3, 2, 3] 和 [3, 2, 3]。

Answer 1

问题出在 else 这个部分，你只是把结果加了一。其实，从 i 开始的所有子数组，只要结束位置 >= j 的话，都会满足条件。所以你应该把 ans += 1 改成 ans += N - j。除此之外，你的算法看起来没什么大问题。

Answer 2

你的代码有几个问题。首先，正如@maraca提到的，当你遇到一个以j结尾的子数组时，你需要计算所有从同一个i开始并以j、j+1，一直到N结束的子数组。其次，你的循环在处理到数组中的最后一个3时就停止了，因为这时j == N。这就导致子数组[2, 3, 3]和[3, 3]没有被计算在内。

下面是一个稍微修改过的代码版本，解决了这些问题：

def countSubarrays(self, nums: list[int], k: int) -> int:
    i = j = 0
    count = 0
    ans = 0 
    N = len(nums)
    max_element = max(nums)
    
    while i <= N - k:
        while count < k and j < N:
            count += nums[j] == max_element
            j += 1
        
        # have we found a valid subarray? if so, every array from j to N will also be one
        if count == k:
            ans += N - j + 1
        
        # increment i and reset count if required
        count -= nums[i] == max_element
        i += 1
    
    return ans