Subarray Sums Divisible by K

Problem

Subarray Sums Divisible by K

Approach

Let there be a subarray (i, j) whose sum is divisible by k and sum(i,j) = sum(0, j) - sum(0, i - 1).

And the sum for any subarray can be re-written as q*k + r, where q is a quotient and r is the reminder.

Thus,

sum(i,j) = sum(0, j) - sum(0, i - 1)
sum(i, j) = (q1 * k + r1) - (q2 * k + r2)
sum(i, j) = (q1 - q2) * k + (r1 - r2)

We can know in order to make any subarray to be divisible by k, it needs to make sure (r1 - r2) can be divisible by k.

And,

r1 = sum(0, j) % k
r2 = sum(0, i - 1) % k

So if any subarry sum(i, j) is divisible by k, we can say r1 = r2.

Code

class Solution {
    public int subarraysDivByK(int[] A, int K) {
        int N = A.length;
        int[] mods = new int[K];
        int curSum = 0;
        for(int i = 0; i < N; i ++){
            curSum += A[i];
            int mod = curSum % K;
            if(mod < 0){
                mod += K;
                mod %= K;
            }
            mods[mod] ++;
        }
        int ans = 0;
        for(int i = 0; i < K; i ++){
            if(mods[i] > 1){
                ans += mods[i] * (mods[i] - 1) / 2;
            }
        }
        // for mod of 0, themselves also should be taken into consideration
        ans += mods[0];
        return ans;
    }
}