Count Subarrays Where Max Element Appears at Least K Times

medium array sliding window two pointers

Problem

Given an integer array nums and a positive integer k, return the number of contiguous subarrays in which the maximum element of the whole array appears at least k times within that subarray.

Inputnums = [1,3,2,3,3], k = 2

Output6

The array max is 3. The subarrays containing 3 at least twice are [1,3,2,3], [1,3,2,3,3], [3,2,3], [3,2,3,3], [2,3,3] and [3,3].

nums (comma-separated)

def countSubarrays(nums, k):
    mx = max(nums)                 # global maximum value
    res = 0
    left = 0
    cnt = 0                        # max-occurrences in window
    for right in range(len(nums)):
        if nums[right] == mx:
            cnt += 1
        while cnt >= k:            # shrink while window has k maxes
            if nums[left] == mx:
                cnt -= 1
            left += 1
        res += left                # left valid starts for this right
    return res

function countSubarrays(nums, k) {
  const mx = Math.max(...nums);      // global maximum value
  let res = 0, left = 0, cnt = 0;    // cnt = maxes in window
  for (let right = 0; right < nums.length; right++) {
    if (nums[right] === mx) cnt++;
    while (cnt >= k) {               // shrink while window has k maxes
      if (nums[left] === mx) cnt--;
      left++;
    }
    res += left;                     // left valid starts for this right
  }
  return res;
}

long countSubarrays(int[] nums, int k) {
    int mx = Integer.MIN_VALUE;
    for (int v : nums) mx = Math.max(mx, v);   // global max
    long res = 0;
    int left = 0, cnt = 0;                      // cnt = maxes in window
    for (int right = 0; right < nums.length; right++) {
        if (nums[right] == mx) cnt++;
        while (cnt >= k) {                      // shrink window
            if (nums[left] == mx) cnt--;
            left++;
        }
        res += left;                            // valid starts for right
    }
    return res;
}

long long countSubarrays(vector<int>& nums, int k) {
    int mx = *max_element(nums.begin(), nums.end()); // global max
    long long res = 0;
    int left = 0, cnt = 0;                           // cnt = maxes in window
    for (int right = 0; right < (int)nums.size(); right++) {
        if (nums[right] == mx) cnt++;
        while (cnt >= k) {                           // shrink window
            if (nums[left] == mx) cnt--;
            left++;
        }
        res += left;                                 // valid starts for right
    }
    return res;
}

Explanation

The only value that matters is the global maximum mx of the array. A subarray is "good" when it contains mx at least k times; every other element is just padding.

We count good subarrays by their right endpoint. Slide right across the array and keep cnt = how many copies of mx sit inside the current window [left, right]. Whenever the window holds k or more maxes, we advance left, dropping elements (and decrementing cnt when the dropped element is a max) until the window has fewer than k maxes.

The key insight is the line res += left. After the while-loop, left is the first index for which the window [left, right] already has fewer than k maxes. That means every start position 0, 1, …, left − 1 gives a subarray ending at right that does contain at least k maxes. So exactly left good subarrays end at this right, and we add that to the answer.

Because left only ever moves forward, each index is visited at most twice (once by right, once by left), giving a single linear pass.

Example: nums = [1,3,2,3,3], k = 2, mx = 3. As right reaches index 3 the window first gathers two 3's; left climbs to 2, adding 2. At index 4 a third 3 arrives, left climbs to 4, adding 4. Total = 2 + 4 = 6.

Time: O(n) Space: O(1)