Count subarrays having sum modulo K same as the length of the subarray

Given an integer K and an array arr[] consisting of N positive integers, the task is to find the number of subarrays whose sum modulo K is equal to the size of the subarray.

Examples:

Input: arr[] = {1, 4, 3, 2}, K = 3
Output: 4
Explanation: 
1 % 3 = 1 
(1 + 4) % 3 = 2 
4 % 3 = 1 
(3 + 2) % 3 = 2 
Therefore, subarrays {1}, {1, 4}, {4}, {3, 2} satisfy the required conditions.

Input: arr[] = {2, 3, 5, 3, 1, 5}, K = 4
Output: 5
Explanation: 
The subarrays (5), (1), (5), (1, 5), (3, 5, 3) satisfy the required condition.

Naive Approach: The simplest approach is to find the prefix sum of the given array, then generate all the subarrays of the prefix sum array and count those subarrays having sum modulo K equal to the length of that subarray. Print the final count of subarrays obtained.



Below is the implementation of the above approach:

C++

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ program of the above approach
 
#include <bits/stdc++.h>
using namespace std;
 
// Function that counts the subarrays
// having sum modulo k equal to the
// length of subarray
 
long long int countSubarrays(
    int a[], int n, int k)
{
 
    // Stores the count
    // of subarrays
    int ans = 0;
 
    // Stores prefix sum
    // of the array
    vector<int> pref;
    pref.push_back(0);
 
    // Calculate prefix sum array
    for (int i = 0; i < n; i++)
        pref.push_back((a[i]
                        + pref[i])
                       % k);
 
    // Generate all the subarrays
    for (int i = 1; i <= n; i++) {
        for (int j = i; j <= n; j++) {
 
            // Check if this subarray is
            // a valid subarray or not
            if ((pref[j] - pref[i - 1] + k) % k
                == j - i + 1) {
                ans++;
            }
        }
    }
 
    // Total count of subarrays
    cout << ans << ' ';
}
 
// Driver Code
int main()
{
    // Given arr[]
    int arr[] = { 2, 3, 5, 3, 1, 5 };
 
    // Size of the array
    int N = sizeof(arr) / sizeof(arr[0]);
 
    // Given K
    int K = 4;
 
    // Function Call
    countSubarrays(arr, N, K);
 
    return 0;
}

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java program for the above approach
import java.util.*;
import java.lang.*;
import java.io.*;
 
class GFG{
     
// Function that counts the subarrays
// having sum modulo k equal to the
// length of subarray
static void countSubarrays(int a[], int n,
                           int k)
{
     
    // Stores the count
    // of subarrays
    int ans = 0;
   
    // Stores prefix sum
    // of the array
    ArrayList<Integer> pref = new ArrayList<>();
    pref.add(0);
   
    // Calculate prefix sum array
    for(int i = 0; i < n; i++)
        pref.add((a[i] + pref.get(i)) % k);
   
    // Generate all the subarrays
    for(int i = 1; i <= n; i++)
    {
        for(int j = i; j <= n; j++)
        {
             
            // Check if this subarray is
            // a valid subarray or not
            if ((pref.get(j) -
                 pref.get(i - 1) + k) %
                     k == j - i + 1)
            {
                ans++;
            }
        }
    }
   
    // Total count of subarrays
    System.out.println(ans);
}
 
// Driver Code
public static void main (String[] args)
throws java.lang.Exception
{
     
    // Given arr[]
    int arr[] = { 2, 3, 5, 3, 1, 5 };
     
    // Size of the array
    int N = arr.length;
     
    // Given K
    int K = 4;
     
    // Function call
    countSubarrays(arr, N, K);
}
}
 
// This code is contributed by bikram2001jha

chevron_right


Python3

filter_none

edit
close

play_arrow

link
brightness_4
code

# Python3 program for the above approach
  
# Function that counts the subarrays
# having sum modulo k equal to the
# length of subarray
def countSubarrays(a, n, k):
  
    # Stores the count
    # of subarrays
    ans = 0
  
    # Stores prefix sum
    # of the array
    pref = []
    pref.append(0)
  
    # Calculate prefix sum array
    for i in range(n):
        pref.append((a[i] + pref[i]) % k)
  
    # Generate all the subarrays
    for i in range(1, n + 1, 1):
        for j in range(i, n + 1, 1):
  
            # Check if this subarray is
            # a valid subarray or not
            if ((pref[j] -
                 pref[i - 1] + k) %
                      k == j - i + 1):
                ans += 1
             
    # Total count of subarrays
    print(ans, end = ' ')
 
# Driver Code
 
# Given arr[]
arr = [ 2, 3, 5, 3, 1, 5 ]
  
# Size of the array
N = len(arr)
  
# Given K
K = 4
  
# Function call
countSubarrays(arr, N, K)
 
# This code is contributed by sanjoy_62

chevron_right


C#

filter_none

edit
close

play_arrow

link
brightness_4
code

// C# program for the above approach 
using System;
using System.Collections.Generic;
 
class GFG{
  
// Function that counts the subarrays
// having sum modulo k equal to the
// length of subarray
static void countSubarrays(int[] a, int n,
                            int k)
{
     
    // Stores the count
    // of subarrays
    int ans = 0;
    
    // Stores prefix sum
    // of the array
    List<int> pref = new List<int>();
    pref.Add(0);
    
    // Calculate prefix sum array
    for(int i = 0; i < n; i++)
        pref.Add((a[i] + pref[i]) % k);
    
    // Generate all the subarrays
    for(int i = 1; i <= n; i++)
    {
        for(int j = i; j <= n; j++)
        {
             
            // Check if this subarray is
            // a valid subarray or not
            if ((pref[j] -
                 pref[i - 1] + k) %
                      k == j - i + 1)
            {
                ans++;
            }
        }
    }
    
    // Total count of subarrays
    Console.WriteLine(ans);
}
  
// Driver Code
public static void Main ()
{
     
    // Given arr[]
    int[] arr = { 2, 3, 5, 3, 1, 5 };
      
    // Size of the array
    int N = arr.Length;
      
    // Given K
    int K = 4;
      
    // Function call
    countSubarrays(arr, N, K);
}
}
 
// This code is contributed by sanjoy_62

chevron_right


Output: 

5







 

Time Complexity: O(N2)
Auxiliary Space: O(1)

Efficient Approach: The idea is to generate the prefix sum of the given array and then the problem reduces to the count of subarray such that (pref[j] – pref[i]) % K equal to (j – i), where j > i and (j − i) ≤ K. Below are the steps:

  • Create an auxiliary array pref[] that stores the prefix sum of the given array.
  • To count the subarray satisfying the above equation, the equation can be written as:

(pref[j] − j) % k = (pref[i] − i) % k, where j > i and (j − i) ≤ K

  • Traverse the prefix array pref[] and for each index j store the count (pref[j] – j) % K in a Map M.
  • For each element pref[i] in the above steps update the count as M[(pref[i] – i % K + K) % K] and increment the frequency of (pref[i] – i % K + K) % K in the Map M.
  • Print the value of the count of subarray after the above steps.

Below is the implementation of the above approach:

C++

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ program of the above approach
 
#include <bits/stdc++.h>
using namespace std;
 
// Function that counts the subarrays
// s.t. sum of elements in the subarray
// modulo k is equal to size of subarray
long long int countSubarrays(
    int a[], int n, int k)
{
    // Stores the count of (pref[i] - i) % k
    unordered_map<int, int> cnt;
 
    // Stores the count of subarray
    long long int ans = 0;
 
    // Stores prefix sum of the array
    vector<int> pref;
    pref.push_back(0);
 
    // Find prefix sum array
    for (int i = 0; i < n; i++)
        pref.push_back((a[i]
                        + pref[i])
                       % k);
 
    // Base Condition
    cnt[0] = 1;
 
    for (int i = 1; i <= n; i++) {
 
        // Remove the index at present
        // after K indices from the
        // current index
        int remIdx = i - k;
 
        if (remIdx >= 0) {
            cnt[(pref[remIdx]
                 - remIdx % k + k)
                % k]--;
        }
 
        // Update the answer for subarrays
        // ending at the i-th index
        ans += cnt[(pref[i]
                    - i % k + k)
                   % k];
 
        // Add the calculated value of
        // current index to count
        cnt[(pref[i] - i % k + k) % k]++;
    }
 
    // Print the count of subarrays
    cout << ans << ' ';
}
 
// Driver Code
int main()
{
    // Given arr[]
    int arr[] = { 2, 3, 5, 3, 1, 5 };
 
    // Size of the array
    int N = sizeof(arr) / sizeof(arr[0]);
 
    // Given K
    int K = 4;
 
    // Function Call
    countSubarrays(arr, N, K);
 
    return 0;
}

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java program for the above approach
import java.util.*;
import java.lang.*;
import java.io.*;
 
class GFG{
     
// Function that counts the subarrays
// having sum modulo k equal to the
// length of subarray
   
static void countSubarrays(int a[], int n,
                                    int k)
{
     
    // Stores the count of (pref[i] - i) % k
    HashMap<Integer, Integer> cnt = new HashMap<>();
   
    // Stores the count of subarray
    long  ans = 0;
     
    // Stores prefix sum of the array
    ArrayList<Integer> pref = new ArrayList<>();
    pref.add(0);
     
    // Find prefix sum array
    for(int i = 0; i < n; i++)
        pref.add((a[i] + pref.get(i)) % k);
   
    // Base Condition
    cnt.put(0, 1);
   
    for(int i = 1; i <= n; i++)
    {
         
        // Remove the index at present
        // after K indices from the
        // current index
        int remIdx = i - k;
   
        if (remIdx >= 0)
        {
            if (cnt.containsKey((pref.get(remIdx) -
                                   remIdx % k + k) % k))
                cnt.put((pref.get(remIdx) -
                          remIdx % k + k) % k,
                cnt.get((pref.get(remIdx) -
                          remIdx % k + k) % k) - 1);
                 
            else
                cnt.put((pref.get(remIdx) -
                          remIdx % k + k) % k, -1);
        }
   
        // Update the answer for subarrays
        // ending at the i-th index
        if (cnt.containsKey((pref.get(i) -
                              i % k + k) % k))
            ans += cnt.get((pref.get(i) -
                             i % k + k) % k);
   
        // Add the calculated value of
        // current index to count
        if (cnt.containsKey((pref.get(i) -
                           i % k + k) % k))
            cnt.put((pref.get(i) -
                      i % k + k) % k,
            cnt.get((pref.get(i) -
                 i % k + k) % k) + 1);
        else
            cnt.put((pref.get(i) -
                      i % k + k) % k, 1);
    }
   
    // Print the count of subarrays
    System.out.println(ans);
}
   
// Driver Code
public static void main (String[] args)
throws java.lang.Exception
{
     
    // Given arr[]
    int arr[] = { 2, 3, 5, 3, 1, 5 };
     
    // Size of the array
    int N = arr.length;
     
    // Given K
    int K = 4;
     
    // Function call
    countSubarrays(arr, N, K);
}
}
 
// This code is contributed by bikram2001jha

chevron_right


Python3

filter_none

edit
close

play_arrow

link
brightness_4
code

# Python3 program of the above approach
 
# Function that counts the subarrays
# s.t. sum of elements in the subarray
# modulo k is equal to size of subarray
def countSubarrays(a, n, k):
 
    # Stores the count of (pref[i] - i) % k
    cnt = {}
  
    # Stores the count of subarray
    ans = 0
  
    # Stores prefix sum of the array
    pref = []
    pref.append(0)
  
    # Find prefix sum array
    for i in range(n):
        pref.append((a[i] + pref[i]) % k)
  
    # Base Condition
    cnt[0] = 1
  
    for i in range(1, n + 1):
     
        # Remove the index at present
        # after K indices from the
        # current index
        remIdx = i - k
  
        if (remIdx >= 0):
            if ((pref[remIdx] -
                   remIdx % k + k) % k in cnt):
                cnt[(pref[remIdx] -
                       remIdx % k + k) % k] -= 1
            else:
                cnt[(pref[remIdx] -
                       remIdx % k + k) % k] = -1
                 
        # Update the answer for subarrays
        # ending at the i-th index
        if (pref[i] - i % k + k) % k in cnt:
            ans += cnt[(pref[i] - i % k + k) % k]
  
        # Add the calculated value of
        # current index to count
        if (pref[i] - i % k + k) % k in cnt:
            cnt[(pref[i] - i % k + k) % k] += 1
        else:
            cnt[(pref[i] - i % k + k) % k] = 1
     
    # Print the count of subarrays
    print(ans, end = ' ')
 
# Driver code 
 
# Given arr[]
arr = [ 2, 3, 5, 3, 1, 5 ]
 
# Size of the array
N = len(arr)
 
# Given K
K = 4
 
# Function call
countSubarrays(arr, N, K)
 
# This code is contributed by divyeshrabadiya07

chevron_right


C#

filter_none

edit
close

play_arrow

link
brightness_4
code

// C# program for
// the above approach
using System;
using System.Collections.Generic;
class GFG{
     
// Function that counts the subarrays
// having sum modulo k equal to the
// length of subarray
static void countSubarrays(int []a, int n,
                           int k)
{
  // Stores the count of
  // (pref[i] - i) % k
  Dictionary<int,
             int> cnt = new Dictionary<int,
                                       int>();
   
  // Stores the count of subarray
  long ans = 0;
 
  // Stores prefix sum of the array
  List<int> pref = new List<int>();
  pref.Add(0);
 
  // Find prefix sum array
  for(int i = 0; i < n; i++)
    pref.Add((a[i] + pref[i]) % k);
 
  // Base Condition
  cnt.Add(0, 1);
 
  for(int i = 1; i <= n; i++)
  {
    // Remove the index at present
    // after K indices from the
    // current index
    int remIdx = i - k;
 
    if (remIdx >= 0)
    {
      if (cnt.ContainsKey((pref[remIdx] -
                           remIdx % k + k) % k))
        cnt[(pref[remIdx] -
             remIdx % k + k) % k] = cnt[(pref[remIdx] -
                                    remIdx % k + k) % k] - 1;
 
      else
        cnt.Add((pref[remIdx] -
                 remIdx % k + k) % k, -1);
    }
 
    // Update the answer for subarrays
    // ending at the i-th index
    if (cnt.ContainsKey((pref[i] -
                         i % k + k) % k))
      ans += cnt[(pref[i] -
                  i % k + k) % k];
 
    // Add the calculated value of
    // current index to count
    if (cnt.ContainsKey((pref[i] -
                         i % k + k) % k))
      cnt[(pref[i] -
           i % k + k) % k] = cnt[(pref[i] -
                                  i % k + k) % k] + 1;
    else
      cnt.Add((pref[i] -
               i % k + k) % k, 1);
  }
 
  // Print the count of subarrays
  Console.WriteLine(ans);
}
   
// Driver Code
public static void Main(String[] args)
 
{
  // Given []arr
  int []arr = {2, 3, 5, 3, 1, 5};
 
  // Size of the array
  int N = arr.Length;
 
  // Given K
  int K = 4;
 
  // Function call
  countSubarrays(arr, N, K);
}
}
 
// This code is contributed by Rajput-Ji

chevron_right


Output: 

5







 

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

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.




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.