Skip to content

2845. Count of Interesting Subarrays 👍

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution {
 public:
  long countInterestingSubarrays(vector<int>& nums, int modulo, int k) {
    long ans = 0;
    int prefix = 0;  // (number of nums[i] % modulo == k so far) % modulo
    unordered_map<int, int> prefixCount{{0, 1}};

    for (const int num : nums) {
      if (num % modulo == k)
        prefix = (prefix + 1) % modulo;
      ans += prefixCount[(prefix - k + modulo) % modulo];
      ++prefixCount[prefix];
    }

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution {
  public long countInterestingSubarrays(List<Integer> nums, int modulo, int k) {
    long ans = 0;
    int prefix = 0; // (number of nums[i] % modulo == k so far) % modulo
    Map<Integer, Integer> prefixCount = new HashMap<>();
    prefixCount.put(0, 1);

    for (final int num : nums) {
      if (num % modulo == k)
        prefix = (prefix + 1) % modulo;
      ans += prefixCount.getOrDefault((prefix - k + modulo) % modulo, 0);
      prefixCount.merge(prefix, 1, Integer::sum);
    }

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Solution:
  def countInterestingSubarrays(self, nums: List[int], modulo: int, k: int) -> int:
    ans = 0
    prefix = 0  # (number of nums[i] % modulo == k so far) % modulo
    prefixCount = collections.Counter({0: 1})

    for num in nums:
      if num % modulo == k:
        prefix = (prefix + 1) % modulo
      ans += prefixCount[(prefix - k + modulo) % modulo]
      prefixCount[prefix] += 1

    return ans