Solution to Leetcode problem 974 Subarray Sums Divisible by K

In this Leetcode problem, we are given an array A and are ask to enter the number of subarrays that have a sum of elements divisible by K .

Problem statement

One obvious solution that is not studied here consists in computing all the sums modulo K and return the number of elements that have the value 0. This solution is O(N*N) if the sums are not entirely re-computed each time.

My solution is alsoO(N*N) in the worst-case. The following solution only optimizes the solution and prevent some wasteful computation in some cases.

Note first that the solution S is the sum of the number of subarrays that start with index i and which sums are 0%K, that we call S[i]. In other words, we divided the problem in N subproblems.

S = sum(S[i]), 0 <= i < A.length

For 0 <= i < A.length, we compute the first element i < j <= A.length so that sum(A(i,j)) = 0 % K, where sum(A(i,j)) is the sum of the subarray A[i,j], j excluded.

If such j doesn’t exist, S[i]=0, otherwise S[i]= 1+S[j].

Using this we can compute all the S[i] for i = A.length — 1 ; i >=0 ; i--, and return the sum of S[i].

class Solution {
    public int subarraysDivByK(int[] A, int K) {
        int L = A.length;
        int[] S = new int[L];
        int sum = 0;
        for (int i = L - 1; i >= 0; i--) {
            int curr = 0;
            for (int j = i; j < L; j++) {
                curr += A[j];
                curr %= K;
                if (curr == 0) {
                    if (j+1 < L) {
                        S[i] = 1 + S[j+1];
                    } else {
                        S[i] = 1;
                    }
                    break;
                }
                S[i] = 0;
            }
            sum += S[i];
        }
        return sum;
    }
}

Thanks for reading! Happy coding :)