Count subarrays such that remainder after dividing sum of elements by K gives count of elements
Last Updated :
03 Jun, 2021
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).
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 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 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# 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
<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)
Explore
Basics & Prerequisites
Data Structures
Algorithms
Advanced
Interview Preparation
Practice Problem