Problem of The Day: Subarray Sums Divisible by K

Problem Statement

Intuition

My first thought to solve this problem was to utilize the prefix sum approach. This technique allows me to keep track of the cumulative sum of the elements in the array, which can be very useful for identifying subarrays that meet certain criteria. The key idea here is to use the modulo operation to determine when a subarray sum is divisible by ( k ).

Approach

To solve this problem, I followed these steps:

1. I initialized a dictionary prefix_sum to keep track of the counts of different modulo values of the cumulative sums.
2. I also initialized curr_sum to store the current cumulative sum and res to store the result, which is the count of subarrays whose sum is divisible by ( k ).
3. I iterated over each element in the array, updating the cumulative sum curr_sum at each step.
4. I calculated the modulo value mod_val of the current cumulative sum with ( k ).
5. If mod_val is zero, it means the subarray from the beginning to the current index is divisible by ( k ), so I incremented the result res.
6. If mod_val has been seen before (exists in prefix_sum), it means there are subarrays ending at the current index that have a sum divisible by ( k ), so I added the count of these subarrays to res.
7. I updated the count of the current mod_val in the prefix_sum dictionary.
8. Finally, I returned the result res.

Complexity

• Time complexity: The time complexity of this approach is ( O(n) ), where ( n ) is the length of the input array. This is because we iterate through the array once.
• Space complexity: The space complexity is ( O(k) ), which is the maximum number of different modulo values we can have. In the worst case, it is ( O(k) ).

Code

class Solution:
    def subarraysDivByK(self, nums: List[int], k: int) -> int:
        n = len(nums)
        prefix_sum = defaultdict(int)
        curr_sum = 0
        res = 0
        for i in range(n):
            curr_sum += nums[i]
            mod_val = curr_sum % k
            if mod_val == 0:
                res += 1
            if mod_val in prefix_sum:
                res += prefix_sum[mod_val]
            prefix_sum[mod_val] += 1
        return res

Editorial

class Solution {
public:
    int subarraysDivByK(vector<int>& nums, int k) {
        int n = nums.size();
        int prefixMod = 0, result = 0;

        // There are k mod groups 0...k-1.
        vector<int> modGroups(k);
        modGroups[0] = 1;

        for (int num : nums) {
            // Take modulo twice to avoid negative remainders.
            prefixMod = (prefixMod + num % k + k) % k;
            // Add the count of subarrays that have the same remainder as the current
            // one to cancel out the remainders.
            result += modGroups[prefixMod];
            modGroups[prefixMod]++;
        }

        return result;
    }
};

• time: O(n + k) where n is the length of nums and k is the given integer
• space: O(k)