## Brute force:
##
## foreach pair of possible subarray start/end indices (l_idx,r_idx): ## O(N^2)
## foreach elem in nums[l_idx:r_idx+1]: ## O(N)
## if elem % m == k:
## elem_count += 1
## if elem_count % m == k:
## interesting_count += 1
## return interesting_count
##
## We are counting and re-counting the same subsets of elems. Use the
## difference of cumsum values instead:
##
## target_elems_to_left = [0]*(1+N_nums)
## for idx,elem in enumerate(nums):
## if elem % m == k:
## target_elems_to_left[1+idx] += 1
## else:
## target_elems_to_left[1+idx] = target_elems_to_left[idx]
##
## foreach pair of possible subarray start/end indices (l_idx,r_idx): ## O(N^2)
## elem_count = target_elems_to_left[1+r_idx] - target_elems_to_left[l_idx] ## O(1)
## if elem_count % m == k:
## interesting_count += 1
## return interesting_count
##
## 0 1 2 3 4 5 6 7 8 idx
## 3 1 4 1 5 9 2 6 5 elems, m=3 k=2
## 0 1 1 1 2 0 2 0 2 elem % m
## 0 0 0 0 1 0 1 0 1 elem % m == k
## 0 0 0 0 1 1 2 2 3 cumsum
## 0 0 0 0 1 1 2 2 0 cumsum % k
## -------------
## -----------
## ---------
## -------
## -----
## ---------------
## -------------
## -----------
## ---------
## -------
## --------
## ------
##
## Some observations:
## (1) If the cumsum up to and including the current elem % m == k,
## then the subarray starting at index 0 and ending with
## the current elem is one we should count.
## (2) For every remainder p, there is another remainder q (the "complement")
## such that (r - q) % m == k.
## (3) Using (2), we can "subtract off" subarrays that have q elems
## to construct new arrays that have % m == k elems
## (4) If we save the counts as we go, we can reuse them in (3).
## (5) Foreach remainder p, we only need the most recent count.
## (6) Instead of searching for the count for p, we can use a map to quicky find it.
## Time: O(n) n = len(nums)
## Space: O(n)
class Solution:
def countInterestingSubarrays(self, nums: List[int], modulo: int, k: int) -> int:
count_by_mod = defaultdict(int)
cs_mod = 0
result = 0
for num in nums:
if num % modulo == k:
cs_mod = (cs_mod + 1) % modulo
if cs_mod == k:
result += 1
cs_comp = (cs_mod - k) % modulo
result += count_by_mod[cs_comp]
count_by_mod[cs_mod] += 1
return result