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++
#include <bits/stdc++.h>
using namespace std;
int sub_arrays( int a[], int n, int k)
{
int sum[n + 2] = { 0 };
for ( int i = 0; i < n; i++) {
a[i]--;
a[i] %= k;
sum[i + 1] += sum[i] + a[i];
sum[i + 1] %= k;
}
int ans = 0, l = 0, r = 0;
map< int , int > mp;
for ( int i = 0; i < n + 1; i++) {
ans += mp[sum[i]];
mp[sum[i]]++;
r++;
if (r - l >= k) {
mp[sum[l]]--;
l++;
}
}
return ans;
}
int main()
{
int a[] = { 1, 4, 2, 3, 5 };
int n = sizeof (a) / sizeof (a[0]);
int k = 4;
cout << sub_arrays(a, n, k);
return 0;
}
|
Java
import java.util.*;
class gfg
{
static int sub_arrays( int []a, int n, int k)
{
int sum[] = new int [n + 2 ] ;
for ( int i = 0 ; i < n+ 2 ; i++)
{
sum[i] = 0 ;
}
for ( int i = 0 ; i < n; i++)
{
a[i]--;
a[i] %= k;
sum[i + 1 ] += sum[i] + a[i];
sum[i + 1 ] %= k;
}
int ans = 0 , l = 0 , r = 0 ;
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++)
{
ans += ( int )mp.get(sum[i]);
temp =( int )mp.get(sum[i]) + 1 ;
mp.put(sum[i], temp);
r++;
if (r - l >= k)
{
temp = ( int )mp.get(sum[l]) - 1 ;
mp.put(sum[l], temp);
l++;
}
}
return ans;
}
public static void main(String args[])
{
int []a = { 1 , 4 , 2 , 3 , 5 };
int n = a.length;
int k = 4 ;
System.out.print(sub_arrays(a, n, k));
}
}
|
Python3
def sub_arrays(a, n, k):
sum = [ 0 for i in range (n + 2 )]
for i in range (n):
a[i] - = 1
a[i] % = k
sum [i + 1 ] + = sum [i] + a[i]
sum [i + 1 ] % = k
ans = 0
l = 0
r = 0
mp = dict ()
for i in range (n + 1 ):
if sum [i] in mp:
ans + = mp[ sum [i]]
mp[ sum [i]] = mp.get( sum [i], 0 ) + 1
r + = 1
if (r - l > = k):
mp[ sum [l]] - = 1
l + = 1
return ans
a = [ 1 , 4 , 2 , 3 , 5 ]
n = len (a)
k = 4
print (sub_arrays(a, n, k))
|
C#
using System;
using System.Collections.Generic;
class gfg
{
static int sub_arrays( int []a, int n, int k)
{
int []sum = new int [n + 2] ;
for ( int i = 0; i < n + 2; i++)
{
sum[i] = 0;
}
for ( int i = 0; i < n; i++)
{
a[i]--;
a[i] %= k;
sum[i + 1] += sum[i] + a[i];
sum[i + 1] %= k;
}
int ans = 0, l = 0, r = 0;
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++)
{
ans += ( int )mp[sum[i]];
temp =( int )mp[sum[i]] + 1;
mp[sum[i]] = temp;
r++;
if (r - l >= k)
{
temp = ( int )mp[sum[l]] - 1;
mp[sum[i]] = temp;
l++;
}
}
return ans;
}
public static void Main(String []args)
{
int []a = { 1, 4, 2, 3, 5 };
int n = a.Length;
int k = 4;
Console.Write(sub_arrays(a, n, k));
}
}
|
Javascript
<script>
function sub_arrays(a, n, k) {
let sum = new Array(n + 2);
for (let i = 0; i < n + 2; i++) {
sum[i] = 0;
}
for (let i = 0; i < n; i++) {
a[i]--;
a[i] %= k;
sum[i + 1] += sum[i] + a[i];
sum[i + 1] %= k;
}
let ans = 0, l = 0, r = 0;
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++) {
ans += mp.get(sum[i]);
temp = mp.get(sum[i]) + 1;
mp.set(sum[i], temp);
r++;
if (r - l >= k) {
temp = mp.get(sum[l]) - 1;
mp.set(sum[i], temp);
l++;
}
}
return ans;
}
let a = [1, 4, 2, 3, 5];
let n = a.length;
let k = 4;
document.write(sub_arrays(a, n, k));
</script>
|
Output:
4
Time Complexity: O(N* log(N))
Auxiliary Space: O(N)
Share your thoughts in the comments
Please Login to comment...