Leetcode 974. Subarray Sums Divisible by K

Problem Explanation

We are provided with an array of integers and we are required to determine the number of contiguous, non-empty subarrays whose sum is divisible by a given integer K.

For example, if we have an array [4,5,0,-2,-3,1] and K = 5. The subarrays with sum divisible by 5 will be 7 which are [4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3].

Problem Solution and Approach

The solution to this problem can be found using prefix sums. A prefix sum array is a data structure that, for a given array of n numbers, tabulates the sum of the first i numbers for 0 <= i < n.

We first initialize the ans variable (which stores our final result) to 0, and the prefix variable (which will store the sum of the first i elements) to 0. We then initialize a count array of size k.

We start iterating through the input array. For every number, we add it to the prefix sum and then take the modulo k of this sum (to prevent overflow). This new prefix value is then used to increase the count of the prefix in the count array.

The answer (ans) is then increased by the count of the prefix. This is done because the sum of all elements from the start of the array up to the current element is stored in the prefix variable, and by adding the count of this prefix value, we account for all such subarrays whose sum is divisible by K.

Python Solution

1
2python
3class Solution:
4    def subarraysDivByK(self, nums, k):
5        ans = prefix = 0
6        count = [1] + [0] * k
7
8        for num in nums:
9            prefix = (prefix + num) % k
10            ans += count[prefix]
11            count[prefix] += 1
12
13        return ans

Java Solution

1
2java
3class Solution {
4    public int subarraysDivByK(int[] nums, int K) {
5        int[] count = new int[K];
6        count[0] = 1;
7        int prefix = 0, res = 0;
8        for (int num : nums) {
9            prefix = (prefix + num % K + K) % K;
10            res += count[prefix]++;
11        }
12        return res;
13    }
14}

JavaScript Solution

1
2javascript
3class Solution {
4    subarraysDivByK(nums, K) {
5        let count = new Array(K).fill(0);
6        count[0] = 1;
7        let prefix = 0, res = 0;
8        for (let num of nums) {
9            prefix = (prefix + num % K + K) % K;
10            res += count[prefix]++;
11        }
12        return res;
13    }
14}

C++ Solution

1
2cpp
3class Solution {
4public:
5    int subarraysDivByK(vector<int>& nums, int K) {
6        vector<int> count(K);
7        count[0] = 1;
8        int prefix = 0, res = 0;
9        for (int num : nums) {
10            prefix = (prefix + num % K + K) % K;
11            res += count[prefix]++;
12        }
13        return res;
14    }
15};

C# Solution

1
2csharp
3public class Solution {
4    public int SubarraysDivByK(int[] nums, int K) {
5        int[] count = new int[K];
6        count[0] = 1;
7        int prefix = 0, res = 0;
8        foreach (int num in nums) {
9            prefix = (prefix + num % K + K) % K;
10            res += count[prefix]++;
11        }
12        return res;
13    }
14}

Conclusion

In summary, we used the prefix sum array to solve this problem, where we iterated over the given array and incremented the count of each prefix sum by 1. To find the total number of subarrays divisible by K, we added the count of each prefix sum to the result.

Python, Java, JavaScript, C++ and C# solutions were provided. Each of them uses a similar strategy to solve the problem.

It is important to note that the prefix sum approach is a common and useful technique in solving array-related problems, especially those involving subarrays. Once you understand this concept, many problems become much easier to solve.

In this particular problem, the key insight was that two subarrays with the same mo of cumulative sum implies that the sum of the elements between these two indices must be divisible by K. By counting the occurrences of all possible mods using count array, we are then able to easily calculate the final answer.

These solutions exhibit a time complexity of O(N) and a space complexity of O(K), where N is the length of the input array and K is the provided integer.


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫