Question

2537. 统计好子数组的数目

English Version

题目描述

给你一个整数数组 nums 和一个整数 k ，请你返回 nums 中 好 子数组的数目。

一个子数组 arr 如果有 至少 k 对下标 (i, j) 满足 i < j 且 arr[i] == arr[j] ，那么称它是一个 好 子数组。

子数组 是原数组中一段连续 非空 的元素序列。

 

示例 1：

输入：nums = [1,1,1,1,1], k = 10
输出：1
解释：唯一的好子数组是这个数组本身。

示例 2：

输入：nums = [3,1,4,3,2,2,4], k = 2
输出：4
解释：总共有 4 个不同的好子数组：
- [3,1,4,3,2,2] 有 2 对。
- [3,1,4,3,2,2,4] 有 3 对。
- [1,4,3,2,2,4] 有 2 对。
- [4,3,2,2,4] 有 2 对。

 

提示：

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

解法

方法一：哈希表 + 双指针

如果一个子数组中包含 $k$ 对相同的元素，那么包含这个子数组的数组一定至少包含 $k$ 对相同的元素。

我们用一个哈希表 $cnt$ 统计窗口内数组元素出现的次数，用 $cur$ 统计窗口内相同元素的对数，用 $i$ 维护窗口的左端点。

遍历数组 $nums$，我们将当前元素 $x$ 作为右端点，那么窗口内相同元素的对数将增加 $cnt[x]$，同时将 $x$ 的出现次数加一，即 $cnt[x] \leftarrow cnt[x] + 1$。接下来，我们循环判断移出左端点后窗口内相同元素的对数是否大于等于 $k$，如果大于等于 $k$，那么我们将左端点元素的出现次数减一，即 $cnt[nums[i]] \leftarrow cnt[nums[i]] - 1$，同时将窗口内相同元素的对数减去 $cnt[nums[i]]$，即 $cur \leftarrow cur - cnt[nums[i]]$，同时将左端点右移，即 $i \leftarrow i + 1$。此时窗口左端点以及左侧的所有元素都可以作为当前右端点的左端点，因此答案加上 $i + 1$。

最后，我们返回答案即可。

时间复杂度 $O(n)$，空间复杂度 $O(n)$。其中 $n$ 为数组 $nums$ 的长度。

class Solution:
  def countGood(self, nums: List[int], k: int) -> int:
    cnt = Counter()
    ans = cur = 0
    i = 0
    for x in nums:
      cur += cnt[x]
      cnt[x] += 1
      while cur - cnt[nums[i]] + 1 >= k:
        cnt[nums[i]] -= 1
        cur -= cnt[nums[i]]
        i += 1
      if cur >= k:
        ans += i + 1
    return ans

class Solution {
  public long countGood(int[] nums, int k) {
    Map<Integer, Integer> cnt = new HashMap<>();
    long ans = 0, cur = 0;
    int i = 0;
    for (int x : nums) {
      cur += cnt.getOrDefault(x, 0);
      cnt.merge(x, 1, Integer::sum);
      while (cur - cnt.get(nums[i]) + 1 >= k) {
        cur -= cnt.merge(nums[i++], -1, Integer::sum);
      }
      if (cur >= k) {
        ans += i + 1;
      }
    }
    return ans;
  }
}

class Solution {
public:
  long long countGood(vector<int>& nums, int k) {
    unordered_map<int, int> cnt;
    long long ans = 0;
    long long cur = 0;
    int i = 0;
    for (int& x : nums) {
      cur += cnt[x]++;
      while (cur - cnt[nums[i]] + 1 >= k) {
        cur -= --cnt[nums[i++]];
      }
      if (cur >= k) {
        ans += i + 1;
      }
    }
    return ans;
  }
};

func countGood(nums []int, k int) int64 {
  cnt := map[int]int{}
  ans, cur := 0, 0
  i := 0
  for _, x := range nums {
    cur += cnt[x]
    cnt[x]++
    for cur-cnt[nums[i]]+1 >= k {
      cnt[nums[i]]--
      cur -= cnt[nums[i]]
      i++
    }
    if cur >= k {
      ans += i + 1
    }
  }
  return int64(ans)
}