Question

2962. 统计最大元素出现至少 K 次的子数组

English Version

题目描述

给你一个整数数组 nums 和一个 正整数 k 。

请你统计有多少满足 「 nums 中的 最大 元素」至少出现 k 次的子数组，并返回满足这一条件的子数组的数目。

子数组是数组中的一个连续元素序列。

 

示例 1：

输入：nums = [1,3,2,3,3], k = 2
输出：6
解释：包含元素 3 至少 2 次的子数组为：[1,3,2,3]、[1,3,2,3,3]、[3,2,3]、[3,2,3,3]、[2,3,3] 和 [3,3] 。

示例 2：

输入：nums = [1,4,2,1], k = 3
输出：0
解释：没有子数组包含元素 4 至少 3 次。

 

提示：

• 1 <= nums.length <= 105
• 1 <= nums[i] <= 106
• 1 <= k <= 105

解法

方法一：双指针

不妨记数组中的最大值为 $mx$。

我们用两个指针 $i$ 和 $j$ 维护一个滑动窗口，使得 $[i, j)$ 之间的子数组中，有 $k$ 个元素等于 $mx$。如果我们固定左端点 $i$，那么所有大于等于 $j-1$ 的右端点都满足条件，一共有 $n - (j - 1)$ 个。

因此，我们枚举左端点 $i$，用指针 $j$ 维护右端点，用一个变量 $cnt$ 记录当前窗口中等于 $mx$ 的元素个数，当 $cnt$ 大于等于 $k$ 时，我们就找到了满足条件的子数组，将答案增加 $n - (j - 1)$。然后我们更新 $cnt$，继续枚举下一个左端点。

时间复杂度 $O(n)$，其中 $n$ 是数组 $nums$ 的长度。空间复杂度 $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;
}