Question

2537. Count the Number of Good Subarrays

中文文档

Description

Given an integer array nums and an integer k, return _the number of good subarrays of_ nums.

A subarray arr is good if it there are at least k pairs of indices (i, j) such that i < j and arr[i] == arr[j].

A subarray is a contiguous non-empty sequence of elements within an array.

 

Example 1:

Input: nums = [1,1,1,1,1], k = 10
Output: 1
Explanation: The only good subarray is the array nums itself.

Example 2:

Input: nums = [3,1,4,3,2,2,4], k = 2
Output: 4
Explanation: There are 4 different good subarrays:
- [3,1,4,3,2,2] that has 2 pairs.
- [3,1,4,3,2,2,4] that has 3 pairs.
- [1,4,3,2,2,4] that has 2 pairs.
- [4,3,2,2,4] that has 2 pairs.

 

Constraints:

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

Solutions

Solution 1: Hash Table + Two Pointers

If a subarray contains $k$ pairs of identical elements, then an array that contains this subarray must contain at least $k$ pairs of identical elements.

We use a hash table $cnt$ to count the number of occurrences of each element in the window, use $cur$ to count the number of pairs of identical elements in the window, and use $i$ to maintain the left endpoint of the window.

We traverse the array $nums$, take the current element $x$ as the right endpoint, then the number of pairs of identical elements in the window will increase by $cnt[x]$, and the occurrence times of $x$ will be increased by one, i.e., $cnt[x] \leftarrow cnt[x] + 1$. Next, we loop to judge whether the number of pairs of identical elements in the window is greater than or equal to $k$ after removing the left endpoint. If it is greater than or equal to $k$, then we decrease the occurrence times of the left endpoint element by one, i.e., $cnt[nums[i]] \leftarrow cnt[nums[i]] - 1$, and decrease the number of pairs of identical elements in the window by $cnt[nums[i]]$, i.e., $cur \leftarrow cur - cnt[nums[i]]$, and move the left endpoint to the right, i.e., $i \leftarrow i + 1$. At this time, all elements to the left of the window left endpoint and the left endpoint itself can be used as the left endpoint of the current right endpoint, so the answer is increased by $i + 1$.

Finally, we return the answer.

The time complexity is $O(n)$, and the space complexity is $O(n)$, where $n$ is the length of the array $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)
}