Count subarrays such that remainder after dividing sum of elements by K gives count of elements

Given an array arr[] of size N and an element K. The task is to find the number of sub-arrays of the given array such that the remainder when dividing the sum of its elements by K is equal to the number of elements in the subarray.

Examples:

Input: arr[] = {1, 4, 2, 3, 5}, K = 4 
Output: 4 
{1}, {1, 4, 2}, {4, 2} and {5} 
are the only valid subarrays.
Input: arr[] = {4, 2, 4, 2, 4, 2, 4, 2}, K = 4 
Output: 7 

Approach: Let’s define a sequence S_n such that S_i = A₁ + A₂ + ··· + A_i and S₀ = 0. Then, the condition that a contiguous subsequence A_i+1, …, A_j is valid can be represented as (S_j – S_i) % K = j – i. 
This equation can then be transformed into the following equivalent conditions: 
(S_j – j) % K = (S_i – i) % K and j – i < K. 
Therefore, for each j(1 ? j ? N), count the number of j – K < i < j such that (S_j – j) % K = (S_i – i) % K. For j the segment needed to be searched is (j – K, j), and for j + 1, it is (j – K + 1, j + 1), and these differ only by one element at the leftmost and rightmost, so in order to search for (j + 1)^th after searching for j^th element, only discard the leftmost element and add the rightmost element. Operations of discarding or adding can be performed quickly by managing the number of S_i – i‘s by using associative arrays (such as map in C++ or dict in Python).

Below is the implementation of the above approach: 

C++

// C++ implementation of the approach
#include <bits/stdc++.h>
using namespace std;

// Function to return the number of subarrays
// of the given array such that the remainder
// when dividing the sum of its elements
// by K is equal to the number of its elements
int sub_arrays(int a[], int n, int k)
{

    // To store prefix sum
    int sum[n + 2] = { 0 };

    for (int i = 0; i < n; i++) {

        // We are dealing with zero
        // indexed array
        a[i]--;

        // Taking modulus value
        a[i] %= k;

        // Prefix sum
        sum[i + 1] += sum[i] + a[i];
        sum[i + 1] %= k;
    }

    // To store the required answer, the left
    // index and the right index
    int ans = 0, l = 0, r = 0;

    // To store si - i value
    map<int, int> mp;

    for (int i = 0; i < n + 1; i++) {

        // Include sum
        ans += mp[sum[i]];
        mp[sum[i]]++;

        // Increment the right index
        r++;

        // If subarray has at least
        // k elements
        if (r - l >= k) {
            mp[sum[l]]--;
            l++;
        }
    }

    // Return the required answer
    return ans;
}

// Driver code
int main()
{
    int a[] = { 1, 4, 2, 3, 5 };
    int n = sizeof(a) / sizeof(a[0]);

    int k = 4;

    // Function call
    cout << sub_arrays(a, n, k);

    return 0;
}

Java

// Java implementation of the approach
import java.util.*;

class gfg
{

    // Function to return the number of subarrays
    // of the given array such that the remainder
    // when dividing the sum of its elements
    // by K is equal to the number of its elements
    static int sub_arrays(int []a, int n, int k)
    {

        // To store prefix sum
        int sum[] = new int[n + 2] ;

        for (int i = 0; i < n+2; i++)
        {
            sum[i] = 0;
        }

        for (int i = 0; i < n; i++)
        {

            // We are dealing with zero
            // indexed array
            a[i]--;

            // Taking modulus value
            a[i] %= k;

            // Prefix sum
            sum[i + 1] += sum[i] + a[i];
            sum[i + 1] %= k;
        }

        // To store the required answer, the left
        // index and the right index
        int ans = 0, l = 0, r = 0;

        // To store si - i value
        HashMap<Integer, Integer> mp = new HashMap<Integer, Integer>();

        for (int i = 0; i < n + 1; i++)
        {
            mp.put(sum[i], 0);
        }
        int temp;

        for (int i = 0; i < n + 1; i++)
        {

            // Include sum
            ans += (int)mp.get(sum[i]);
            temp =(int)mp.get(sum[i]) + 1;
            mp.put(sum[i], temp);

            // Increment the right index
            r++;

            // If subarray has at least
            // k elements
            if (r - l >= k)
            {
                //mp[sum[l]]--;
                temp = (int)mp.get(sum[l]) - 1;
                mp.put(sum[l], temp);
                l++;
            }
        }

        // Return the required answer
        return ans;
    }

    // Driver code
    public static void main(String args[])
    {
        int []a = { 1, 4, 2, 3, 5 };

        int n = a.length;

        int k = 4;

        // Function call
        System.out.print(sub_arrays(a, n, k));

    }
}

// This code is contributed by AnkitRai01

Python3

# Python3 implementation of the approach

# Function to return the number of
# subarrays of the given array
# such that the remainder when dividing
# the sum of its elements by K is
# equal to the number of its elements
def sub_arrays(a, n, k):

    # To store prefix sum
    sum = [0 for i in range(n + 2)]

    for i in range(n):

        # We are dealing with zero
        # indexed array
        a[i] -= 1

        # Taking modulus value
        a[i] %= k

        # Prefix sum
        sum[i + 1] += sum[i] + a[i]
        sum[i + 1] %= k

    # To store the required answer,
    # the left index and the right index
    ans = 0
    l = 0
    r = 0

    # To store si - i value
    mp = dict()

    for i in range(n + 1):

        # Include sum
        if sum[i] in mp:
            ans += mp[sum[i]]
        mp[sum[i]] = mp.get(sum[i], 0) + 1

        # Increment the right index
        r += 1

        # If subarray has at least
        # k elements
        if (r - l >= k):
            mp[sum[l]] -= 1
            l += 1

    # Return the required answer
    return ans

# Driver code
a = [1, 4, 2, 3, 5]
n = len(a)

k = 4

# Function call
print(sub_arrays(a, n, k))

# This code is contributed by Mohit Kumar

C#

// C# implementation of the approach
using System;
using System.Collections.Generic;

class gfg
{
    // Function to return the number of subarrays
    // of the given array such that the remainder
    // when dividing the sum of its elements
    // by K is equal to the number of its elements
    static int sub_arrays(int []a, int n, int k)
    {

        // To store prefix sum
        int []sum = new int[n + 2] ;

        for (int i = 0; i < n + 2; i++)
        {
            sum[i] = 0;
        }

        for (int i = 0; i < n; i++)
        {

            // We are dealing with zero
            // indexed array
            a[i]--;

            // Taking modulus value
            a[i] %= k;

            // Prefix sum
            sum[i + 1] += sum[i] + a[i];
            sum[i + 1] %= k;
        }

        // To store the required answer, the left
        // index and the right index
        int ans = 0, l = 0, r = 0;

        // To store si - i value
        Dictionary<int, int> mp = new Dictionary<int, int>();

        for (int i = 0; i < n + 1; i++)
        {
            if(!mp.ContainsKey(sum[i]))
                mp.Add(sum[i], 0);
        }
        int temp;

        for (int i = 0; i < n + 1; i++)
        {

            // Include sum
            ans += (int)mp[sum[i]];
            temp =(int)mp[sum[i]] + 1;
            mp[sum[i]] = temp;

            // Increment the right index
            r++;

            // If subarray has at least
            // k elements
            if (r - l >= k)
            {
                //mp[sum[l]]--;
                temp = (int)mp[sum[l]] - 1;
                mp[sum[i]] = temp;
                l++;
            }
        }

        // Return the required answer
        return ans;
    }

    // Driver code
    public static void Main(String []args)
    {
        int []a = { 1, 4, 2, 3, 5 };

        int n = a.Length;

        int k = 4;

        // Function call
        Console.Write(sub_arrays(a, n, k));
    }
}

// This code is contributed by 29AjayKumar

Javascript

<script>

// JavaScript implementation of the approach

// Function to return the number of subarrays
// of the given array such that the remainder
// when dividing the sum of its elements
// by K is equal to the number of its elements
function sub_arrays(a, n, k) {

    // To store prefix sum
    let sum = new Array(n + 2);

    for (let i = 0; i < n + 2; i++) {
        sum[i] = 0;
    }

    for (let i = 0; i < n; i++) {

        // We are dealing with zero
        // indexed array
        a[i]--;

        // Taking modulus value
        a[i] %= k;

        // Prefix sum
        sum[i + 1] += sum[i] + a[i];
        sum[i + 1] %= k;
    }

    // To store the required answer, the left
    // index and the right index
    let ans = 0, l = 0, r = 0;

    // To store si - i value
    let mp = new Map();

    for (let i = 0; i < n + 1; i++) {
        if (!mp.has(sum[i]))
            mp.set(sum[i], 0);
    }
    let temp;

    for (let i = 0; i < n + 1; i++) {

        // Include sum
        ans += mp.get(sum[i]);
        temp = mp.get(sum[i]) + 1;
        mp.set(sum[i], temp);

        // Increment the right index
        r++;

        // If subarray has at least
        // k elements
        if (r - l >= k) {
            //mp[sum[l]]--;
            temp = mp.get(sum[l]) - 1;
            mp.set(sum[i], temp);
            l++;
        }
    }

    // Return the required answer
    return ans;
}

// Driver code

let a = [1, 4, 2, 3, 5];

let n = a.length;

let k = 4;

// Function call
document.write(sub_arrays(a, n, k));

// This code is contributed by _saurabh_jaiswal

</script>

Output:

4

Time Complexity: O(N* log(N))
Auxiliary Space: O(N)