PROBLEM LINK: Return of Karunee
DIFFICULTY:
Medium
PREREQUISITES:
Maths, Combinatorics, Prefix sum
PROBLEM:
Given an array of digits from 0 - 9. You have to count total number of evil subarrays of the given array. A subarray is evil if sum of all elements in that subarray is equal to the size of the subarray.
QUICK EXPLANATION:
For subtask 1 we can easily do brute force using two nested loops and overall time complexity would be O(T*M^2).
For subtask 2, Store the frequency of A_i - i for ( 0 \leq i < M) . Each distinct key will contribute \binom{n}{2} = \frac{n*(n-1)}{2} to the total count where n is the frequency of the key.
EXPLANATION:
First of all we precalculate a prefix sum array P, where P_i = \displaystyle\sum_{j = 0}^{i-1} A_i . It means P_x is sum of first x elements of array A.
Let’s assume [l,r] is an evil subarray, then
P_r - P_{l-1} = r - l + 1
\Rightarrow P_r - P_{l-1} = r - (l - 1)
\Rightarrow P_r - r = P_{l-1} - (l-1)
From here we can say that if there are K such i whose P_i - i value is same then selecting any two i from it, always results in an evil subarray. So we just need to store the frequency of all the distinct values of P_i - i ( can be stored using a hash map ) and for each frequency x we have to add \frac{x*(x-1)}{2} to the answer.
Overall time - complexity of this approach is O(T*M).
SOLUTIONS:
For Subtask #1:
Editorialist's Solution
#include <bits/stdc++.h>
using namespace std;
#define ll long long int
void solve()
{
ll t ;
cin >> t ;
while (t--)
{
ll n ;
cin >> n ;
string s ;
cin >> s ;
ll ans = 0 ;
for (ll i = 0 ; i < n ; i++)
{
ll sum = 0 ;
for (ll j = i ; j < n ; j++)
{
sum += s[j] - '0' ;
if (sum == (j - i + 1 ))
ans ++ ;
}
}
cout << ans << "\n" ;
}
}
int main()
{
solve();
return 0;
}
For Subtask #2:
Editorialist's Solution
#include <bits/stdc++.h>
using namespace std;
#define ll long long int
void solve()
{
ll t ;
cin >> t ;
while (t--)
{
ll n ;
cin >> n ;
string s ;
cin >> s ;
ll p[n + 1 ] ;
unordered_map <ll , ll> mp ;
ll ans = 0 , sum = 0 ;
// Calculating Prefix sum array P
p[0] = 0 ;
for (ll i = 0 ; i < n ; i++)
{
p[i + 1] = p[i] + (s[i] - '0') ;
}
// Storing frequency of distict value of Pi - i
for (ll i = 0 ; i <= n ; i++)
{
mp[p[i] - i]++ ;
}
// Adding (x*(x-1))/2 for each frequency x
for (auto &i : mp)
{
ans += ((i.second) * (i.second - 1 )) / 2 ;
}
cout << ans << "\n" ;
}
}
int main()
{
solve();
return 0;
}