As you are only interested in numbers divisible by K, you can do all computations modulo K.
Consider the cumulative sum array S such that S[i] = S[0] + S[1] + ... + S[i]
. Then (P, Q) is a slice divisible by K iff S[P] = S[Q]
(remember we do all computations modulo K). So you just have to count for each possible value of [0 ,..., K-1] how many times it appears in S.
Here is some pseudocode:
B = new array( K )
B[0]++
s = 0
for i = 0 to N - 1
s = ( s + A[i] ) % K
B[s]++
ans = 0
for i = 0 to K - 1
ans = ans + B[i] * ( B[i] - 1 ) / 2
Once you know that they are x cells in S that have value i, you want to count the number of slices the start in a cell with value i and ends in a cell with value i, this number is x ( x - 1 ) / 2
. To solve edge problems, we add one cell with value 0.
What does x ( x - 1 ) / 2
stands for: Let's assume our array is [4, 5, 0] and frequency of 4 as prefix sum is x, which is 3 in this case. Now we can conclude from value of x, that there are at least x-1 numbers which are either divisible by k or have mod k equals to 0. Now total possible sub arrays out of these x-1 numbers are 1 + 2 + 3 ... + ( x - 1 ) which is ( ( x - 1 ) * ( ( x - 1 ) + 1 ) / 2
. (Standard formula for summation from 1 to N where N stands for ( x - 1 ).