Problem of The Day: Count of Interesting Subarrays

Problem Statement

leetcode problem link

Brute Force [TLE]

class Solution:
    def countInterestingSubarrays(self, nums: List[int], modulo: int, k: int) -> int:
        res = 0
        N = len(nums)
        arr = []
        for i in range(N):
            arr.append(nums[i] % modulo)

        for i in range(N):
            cnt = 0
            for j in range(i, N):
                if arr[j] == k:
                    cnt += 1
                if cnt % modulo == k:
                    res += 1
        return res

Editorial

Approach: Prefix Sum

class Solution:
    def countInterestingSubarrays(
        self, nums: List[int], modulo: int, k: int
    ) -> int:
        n = len(nums)
        cnt = Counter([0])
        res = 0
        prefix = 0
        for i in range(n):
            prefix += 1 if nums[i] % modulo == k else 0
            res += cnt[(prefix - k + modulo) % modulo]
            cnt[prefix % modulo] += 1
        return res