简体   繁体   中英

Number of subarrays divisible by k

I had the following question in an interview and, in spite of the fact that I gave a working implementation, it wasn't efficient enough.

A slice of array A is any pair of integers (P, Q) such that 0 ≤ P ≤ Q < N. A slice (P, Q) of array A is divisible by K if the number A[P] + A[P+1] + ... + A[Q−1] + A[Q] is divisible by K.

The function I was asked to write, had to return the number of slices divisible by K. The expected time complexity was O(max(N, K)) and space complexity was O(K).

My solution was the simplest, one loop inside another and check every slice: O(n^2)

I have been thinking but I really can't figure out how can I do it in O(max(N, K)).

It may be a variant of the subset sum problem , but I don't know how to count every subarray.

EDIT: Elements in array could be negatives. Here is an example:

A = {4, 5, 0, -2, -3, 1}, K = 5

Function must return 7, because there are 7 subarrays which sums are divisible by 5
{4, 5, 0, -2, -3, 1}
{5}
{5, 0}
{5, 0, -2, -3}
{0}
{0, -2, -3}
{-2, -3}

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 ).

For a given number X ...

The basic idea:

the sum from the first element to b = the sum from the first element to a
                                    + the sum of the elements between the two

So:

the sum of the elements between the two = the sum from the first element to b
                                        - the sum from the first element to a

Then, if those sums on the right both have the same remainder when divided by X , the remainders will cancel out and sum of the elements between the two will be divisible by X . An elaboration:

C = the sum of the elements between the two
B = the sum from the first element to b
A = the sum from the first element to a

Now we can convert B to the form PX + Q and A to the form RX + S , for some integers P , Q , R and S , with 0 <= Q, S < X . Here, by definition, Q and S would be the respective remainders of B and A being divided by X .

Then we have:

C = (PX + Q) - (RX + S)
C = PX + Q - RX - S
C = PX - RX + Q - S
C = (P-R)X + Q - S

Clearly (PR)X is divisible by X (the result is simply (PR) ). Now we just need Q - S to be divisible by X , but, since 0 <= Q, S < X , they'll need to be equal.

Example:

Let B = 13 , A = 7 , X = 3 .

Here B % X = 1 and A % X = 1 .

We can rewrite B as 4*3 + 1 and A as 2*3 + 1 .

Then C = 4*3 + 1 - 2*3 - 1 = 2*3 , which is divisible by 3 .

High-level approach:

Construct a hash-map which will store the cumulative sum of all the numbers thus far mod X mapped to the count of how often that remainder value appears (constructed in expected O(n) ).

Increase 0 's value by one - this corresponds to the start of the array.

Initialize a count to 0.

Go through the hash-map and add nC2 (= value!/(2*(value-2)!) ) to the count. The 2 we're choosing here are the starting and ending positions of the subarray.

The count is the desired value.

Running time:

Expected O(n) .

Example:

Input:    0  5  3  8  2  1
X = 3

Sum:   0  0  5  8 16 18 19
Mod 3: 0  0  2  2  1  0  1

Map:
  0 -> 3
  2 -> 2
  1 -> 2

Count = 3! / 2*(3-2)! = 3  +
        2! / 2*(2-2)! = 1  +
        2! / 2*(2-2)! = 1
      = 5

The subarrays will be:

0  5  3  8  2  1
-                     0                 =  0 % 3 = 0
-------------         0 + 5 + 3 + 8 + 2 = 18 % 3 = 0
   ----------         5 + 3 + 8 + 2     = 18 % 3 = 0
      -               3                 =  3 % 3 = 0
            ----      2 + 1             =  3 % 3 = 0
    private int GetSubArraysCount(int[] A, int K)
    {
        int N = A.Length;
        int[] B = new int[K];
        for (int i = 0; i < B.Length; i++)
        {
            B[i] = 0;
        }
        B[0]++;
        int s = 0;
        for (int i = 0; i < N; i++)
        {
            s = (s + A[i]) % K;
            while (s < 0)
            {
                s += K;
            }
            B[s]++;
        }
        int ans = 0;
        for (int i = 0; i <= K - 1; i++)
        {
            ans += B[i] * (B[i] - 1) / 2;
        }
        return ans;
    }
static void Main(string[] args)
    {
        int[] A = new int[] { 4, 5, 0, -2, -3, 1 };
        int sum = 0;
        int i, j;
        int count = 0;
        for (i = 0; i < A.Length; i++)
        {
            for (j = 0; j < A.Length; j++)
            {
                if (j + i < 6)
                    sum += A[j + i];
                if ((sum % 5) == 0)
                    count++;

            }
            sum = 0;
        }
        Console.WriteLine(count);
        Console.ReadLine();


    }

Here is a Java implementation of the solution proposed by @Thomash.

The second loop is not necessary, because we can directly increase the answer by the current value and then increment it.

To avoid negative array index, we also have to adjust module computation.

public static int countSubarrays(int[] nums, int k) {
    int[] cache = new int[k];
    cache[0]++;
    int s = 0, counter = 0;
    for (int i = 0; i < nums.length; i++) {
        s = ((s + nums[i]) % k + k) % k;
        counter += cache[s];
        cache[s]++;
    }

    return counter;
}

Example :-

Input Array

int [] nums = {4,3,1,2,1,5,2};

K is 3

Consecutive sum

4,7,8,10,11,16,18

Divide above consecutive sum array by 3

1,1,2,1,2,1,0

so we have got four 1's, Two 2's, one 0's

so total count will be (4*3)/2 + (2*1)/2 + (2*1)/2 = 8

(4*3)/2 comes from select any two 1's out of four ie nC2 = n(n-1)/2

Here is the program

public static long countSubArrayDivByK(int k, int[] nums) {

    Map<Integer, Integer> modulusCountMap = new HashMap<Integer, Integer>();
    int [] consecSum = new int[nums.length];
    consecSum[0]=nums[0];

    for(int i=1;i<nums.length;i++){
        consecSum[i]= consecSum[i-1] +nums[i];
    }

    for(int i=0;i<nums.length;i++){
        consecSum[i]= consecSum[i]%k;

            if(consecSum[i]==0 && modulusCountMap.get(consecSum[i])==null){
                modulusCountMap.put(consecSum[i], 2);
            }else{
                modulusCountMap.put(consecSum[i], modulusCountMap.get(consecSum[i])==null ? 1 : modulusCountMap.get(consecSum[i])+1);
            }

    }

    int count = 0;

    for (Integer val : modulusCountMap.values()) {
        count = count +  (val*(val-1))/2;
    }

    return count;
}

Optimized version of above

static long customOptimizedCountSubArrayDivByK(int k, int[] nums) {

        Map<Integer, Integer> modulusCountMap = new HashMap<Integer, Integer>();
        int [] quotient = new int[nums.length];
        quotient[0]=nums[0]%3;



        if(quotient[0]==0){
            modulusCountMap.put(quotient[0], 2);
        }else{
            modulusCountMap.put(quotient[0], 1);
        }


        for(int i=1;i<nums.length;i++){
            quotient[i]= (quotient[i-1] + nums[i])%3;


                if(quotient[i]==0 && modulusCountMap.get(quotient[i])==null){
                    modulusCountMap.put(quotient[i], 2);
                }else{
                    modulusCountMap.put(quotient[i], modulusCountMap.get(quotient[i])==null ? 1 : modulusCountMap.get(quotient[i])+1);
                }

        }

        int count = 0;

        for (Integer val : modulusCountMap.values()) {
            count = count +  (val*(val-1))/2;
        }

        return count;
    }

Thanks for your solution , @damluar, it's very neat though! I just want to add some comments.

  1. The output should be 7, not 6 as your output. Because we have 7 subarrays that are divisible by k as below, adding res += storedArray[0]; to fix it.

{4, 5, 0, -2, -3, 1}; {5}; {5, 0}; {5, 0, -2, -3}; {0}; {0, -2, -3}; {-2, -3}

Ref link

  1. Initializing cache[0]++; depends on the language, if using C++, it's needed, but it's unnecessary for java [ link ].

Code:

public class HelloWorld{

public static void main(String []args){
    int [] A = new int[] {4,5,0,-2,-3,1};
    int k = 5;
    int ans=0;
    System.out.println(countSubArray(A, k)); // output = 7

}

public static int countSubArray(int [] nums, int k){
    int [] storedArray = new int[k];
    int sum=0, res=0;
    for(int i=0; i<nums.length; i++){
        sum = (((sum + nums[i]) % k) + k) % k;
        res += storedArray[sum];
        storedArray[sum]++;

    }
    res += storedArray[0];
    return res; 
}
}
public class SubArrayDivisible 
{

    public static void main(String[] args) 
    {
        int[] A = {4, 5, 0, -2, -3, 1};
        SubArrayDivisible obj = new SubArrayDivisible();
        obj.getSubArrays(A,5);
    }

    private void getSubArrays(int[] A,int K)
    {
        int count = 0,s=0;
        for(int i=0;i<A.length;i++)
        {
            s = 0;
            for(int j = i;j<A.length;j++)
            {
                s = s+A[j];
                if((s%K) == 0)
                {
                    System.out.println("Value of S "+s);
                    count++;
                }

            }
        }
        System.out.println("Num of Sub-Array "+count);
    }
}

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM