Count Subarrays Where Max Element Appears at Least K Times sliding window

Question

Problem

Count Subarrays Where Max Element Appears at Least K Times

You are given an integer array nums and a positive integer k. Return the number of subarrays where the maximum element of nums appears at least k times in that subarray. A subarray is a contiguous sequence of elements within an array.

Logic

• Go through the windows, if currently you have less than k maxElement counts keep widening the window
• If you have more than equal to k maxElement counts, increment ans and work on shortening/sliding the window

Code

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
            

Issue

Possible issues could be

• I am failing to consider last subarray that could be shortened or atleast consider to be added in the ans variable.
• I am not sliding the window correctly because I could go right as well as left. Consider nums = [1,3,2,3,3] and K current subarray being [3,2,3] I already satisfied count >= k so instead of considering [3,2,3,3] I would work on shortening the subarray.
• I know that I can brute force a while loop (while count>=k increment i until this agreement gets broken) but I believe this still doesn't solve the problem. This could still happen in an if condition (if count >=k and freezing j until the agreement gets broken) so, any help is appreciated. I wish we can find a solution where if count >=k we dont increment j until we bring back count < k and continue increasing j later
• If you need a specific case, this code fails at nums=[1,3,2,3,3] and k = 2 Output = 2 Expected Output = 6 The two subarrays being [1, 3, 2, 3] and [3, 2, 3]

Answer 1

The problem is in the else clause, you just increment the result. But actually all sub-arrays starting at i with an end index >= j will satisfy the condition. So instead of ans += 1 you have to do ans += N - j. Other than that your algorithm seems fine.

Answer 2

You have a couple of issues with your code. Firstly, as also mentioned by @maraca, when you encounter a sub-array ending at j, you need to count all sub-arrays that start at the same i and end at j, j+1, ..., N. Secondly, your loop terminates the first time you process the final 3 in the array because j == N. This means the sub-arrays [2, 3, 3] and [3, 3] are not counted.

Here's a slightly modified version of your code that doesn't have those issues:

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