Count subarrays divisible by K

Question

Given a sequence of n positive integers we need to count consecutive sub-sequences whose sum is divisible by k.

Constraints : N is up to 10^6 and each element up to 10^9 and K is up to 100

EXAMPLE : Let N=5 and K=3 and array be 1 2 3 4 1

Here answer is 4

Explanation : there exists, 4 sub-sequences whose sum is divisible by 3, they are

3
1 2
1 2 3
2 3 4

My Attempt :

long long int count=0;
for(int i=0;i<n;i++){
    long long int sum=0;
    for(int j=i;j<n;j++)
    {
        sum=sum+arr[j];
        if(sum%k==0)
        {
            count++;
        }
    }
}

But obviously its poor approach. Can their be better approach for this question? Please help.

Complete Question: https://www.hackerrank.com/contests/w6/challenges/consecutive-subsequences

Answer 1

Here is a fast O(n + k) solution:

1)Lets compute prefix sums pref[i](for 0 <= i < n).

2)Now we can compute count[i] - the number of prefixes with sum i modulo k(0 <= i < k). This can be done by iterating over all the prefixes and making count[pref[i] % k]++. Initially, count[0] = 1(an empty prefix has sum 0) and 0 for i != 0.

3)The answer is sum count[i] * (count[i] - 1) / 2 for all i.

4)It is better to compute prefix sums modulo k to avoid overflow.

Why does it work? Let's take a closer a look at a subarray divisible by k. Let's say that it starts in L position and ends in R position. It is divisible by k if and only if pref[L - 1] == pref[R] (modulo k) because their differnce is zero modulo k(by definition of divisibility). So for each fixed modulo, we can pick any two prefixes with this prefix sum modulo k(and there are exactly count[i] * (count[i] - 1) / 2 ways to do it).

Here is my code:

long long get_count(const vector<int>& vec, int k) {
  //Initialize count array.
  vector<int> cnt_mod(k, 0);
  cnt_mod[0] = 1;
  int pref_sum = 0;
  //Iterate over the input sequence.
  for (int elem : vec) {
    pref_sum += elem;
    pref_sum %= k;
    cnt_mod[pref_sum]++;
  }
  //Compute the answer.
  long long res = 0;
  for (int mod = 0; mod < k; mod++)
    res += (long long)cnt_mod[mod] * (cnt_mod[mod] - 1) / 2;
  return res;
}