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++ 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;
}
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++ 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;
}
Output:
5
Time Complexity: O(N)
Auxiliary Space: O(N)
Comments
Post a Comment