leetcode - 2962. Count Subarrays Where Max Element Appears at Least K Times

Description

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.

Example 1:

Input: nums = [1,3,2,3,3], k = 2
Output: 6
Explanation: The subarrays that contain the element 3 at least 2 times are: [1,3,2,3], [1,3,2,3,3], [3,2,3], [3,2,3,3], [2,3,3] and [3,3].

Example 2:

Input: nums = [1,4,2,1], k = 3
Output: 0
Explanation: No subarray contains the element 4 at least 3 times.

Constraints:

1 <= nums.length <= 10^5
1 <= nums[i] <= 10^6
1 <= k <= 10^5

Solution

Solved after help…

Sliding window, expand the window until there are more than k maximum elements inside the window. And the number of subarrays will be: len(nums) - right.
Because if [left, right] is a valid solution, then [left, right + 1], [left, right + 2], ... are also valid solutions. So think about subarrays start from the end of nums, and end with right, the number of subarrays will be: len(nums) - right.

Time complexity: $o(n)$
Space complexity: $o(n)$

Or another way, for the window [left, right], any subarray that starts before left and ends with right are valid solutions. So we can also calculate the number of subarrays by: res += left.

Code

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        left = 0
        cnt = 0
        max_num = max(nums)
        res = 0
        for right in range(len(nums)):
            if nums[right] == max_num:
                cnt += 1
            while cnt >= k:
                res += len(nums) - right
                if nums[left] == max_num:
                    cnt -= 1
                left += 1
        return res

Another way:

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        cnt = 0
        left = 0
        res = 0
        max_ele = max(nums)
        for right in range(len(nums)):
            if nums[right] == max_ele:
                cnt += 1
            while cnt >= k:
                if nums[left] == max_ele:
                    cnt -= 1
                left += 1
            res += left
        return res