Question

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 <= 105
• 1 <= nums[i] <= 106
• 1 <= k <= 105

Solutions

Solution 1: Two Pointers

Let's denote the maximum value in the array as $mx$.

We use two pointers $i$ and $j$ to maintain a sliding window, such that in the subarray between $[i, j)$, there are $k$ elements equal to $mx$. If we fix the left endpoint $i$, then all right endpoints greater than or equal to $j-1$ meet the condition, totaling $n - (j - 1)$.

Therefore, we enumerate the left endpoint $i$, use the pointer $j$ to maintain the right endpoint, use a variable $cnt$ to record the number of elements equal to $mx$ in the current window. When $cnt$ is greater than or equal to $k$, we have found a subarray that meets the condition, and we increase the answer by $n - (j - 1)$. Then we update $cnt$ and continue to enumerate the next left endpoint.

The time complexity is $O(n)$, where $n$ is the length of the array $nums$. The space complexity is $O(1)$.

class Solution:
  def countSubarrays(self, nums: List[int], k: int) -> int:
    mx = max(nums)
    n = len(nums)
    ans = cnt = j = 0
    for x in nums:
      while j < n and cnt < k:
        cnt += nums[j] == mx
        j += 1
      if cnt < k:
        break
      ans += n - j + 1
      cnt -= x == mx
    return ans

class Solution {
  public long countSubarrays(int[] nums, int k) {
    int mx = Arrays.stream(nums).max().getAsInt();
    int n = nums.length;
    long ans = 0;
    int cnt = 0, j = 0;
    for (int x : nums) {
      while (j < n && cnt < k) {
        cnt += nums[j++] == mx ? 1 : 0;
      }
      if (cnt < k) {
        break;
      }
      ans += n - j + 1;
      cnt -= x == mx ? 1 : 0;
    }
    return ans;
  }
}

class Solution {
public:
  long long countSubarrays(vector<int>& nums, int k) {
    int mx = *max_element(nums.begin(), nums.end());
    int n = nums.size();
    long long ans = 0;
    int cnt = 0, j = 0;
    for (int x : nums) {
      while (j < n && cnt < k) {
        cnt += nums[j++] == mx;
      }
      if (cnt < k) {
        break;
      }
      ans += n - j + 1;
      cnt -= x == mx;
    }
    return ans;
  }
};

func countSubarrays(nums []int, k int) (ans int64) {
  mx := slices.Max(nums)
  n := len(nums)
  cnt, j := 0, 0
  for _, x := range nums {
    for ; j < n && cnt < k; j++ {
      if nums[j] == mx {
        cnt++
      }
    }
    if cnt < k {
      break
    }
    ans += int64(n - j + 1)
    if x == mx {
      cnt--
    }
  }
  return
}

function countSubarrays(nums: number[], k: number): number {
  const mx = Math.max(...nums);
  const n = nums.length;
  let [cnt, j] = [0, 0];
  let ans = 0;
  for (const x of nums) {
    for (; j < n && cnt < k; ++j) {
      cnt += nums[j] === mx ? 1 : 0;
    }
    if (cnt < k) {
      break;
    }
    ans += n - j + 1;
    cnt -= x === mx ? 1 : 0;
  }
  return ans;
}