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 Sn such that Si = A1 + A2 + ··· + Ai and S0 = 0. Then, the condition that a contiguous subsequence Ai+1, …, Aj is valid can be represented as (Sj – Si) % K = j – i.
This equation can then be transformed into the following equivalent conditions:
(Sj – j) % K = (Si – i) % K and j – i < K.
Therefore, for each j(1 ≤ j ≤ N), count the number of j – K < i < j such that (Sj – j) % K = (Si – 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 jth 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 Si – i‘s by using associative arrays (such as map in C++ or dict in Python). The total time complexity is about O(NlogK).



Below is the implementation of the above approach:

C++

filter_none

edit
close

play_arrow

link
brightness_4
code

// 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;
}

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

// 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

chevron_right


Python3

filter_none

edit
close

play_arrow

link
brightness_4
code

# 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

chevron_right


C#

filter_none

edit
close

play_arrow

link
brightness_4
code

// 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

chevron_right


Output:

4

Don’t stop now and take your learning to the next level. Learn all the important concepts of Data Structures and Algorithms with the help of the most trusted course: DSA Self Paced. Become industry ready at a student-friendly price.




My Personal Notes arrow_drop_up

Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.