Count subarrays whose sum equals their length

Given an array of integers between 0 to 9 (both inclusive ) of length, N finds the no of subarrays so that the length of the subarray is equal to the sum of the subarray.

Examples :

Input : arr[]: {1, 1, 1}
Output : 6
Explanation :
the no of subarrays are { {1}, {1}, {1}, {1, 1}, {1, 1}, {1, 1, 1} }
and all the subarrays sum is equal to the the length of the subarray.
so the total subarrays are 6.

Input : arr[]: {0, 1, 1, 2}
Output : 4
Explanation :
the subarray which satisfy the given condition is
{{1},{1},{1, 1},{0,1,1,2}}
therefor there are four subarrays that satisfy the given condition.

Approach :

 this question can be done by the greedy approach.

1.  the subarray [l, r] is valid when the sum[0…R] – sum[0…L]  = R – L + 1 
2. therefor  sum[0…R]  – R = sum[0…L]  – L
3. to count this value we will make a map to store the sum[0….i] – I
4. for all the index i the no of subarray is equal to the value map[sum[0….i] – i]
5. repeat this from 0 to N

Time Complexity : O(N log N)

Auxiliary Space : O(N)

Implementation :

#include<bits/stdc++.h> 
#define ll long long 
using namespace std;

ll findSubarrays(vector<int> &arr){
	
	// length of the array
	int n = arr.size(); 
	
	// initialise the 
	ll sum = 0, ans = 0; 
	
	//creating map to store the 
	// occurence 
	map<ll, ll> mp; 
	mp[0] = 1; 
	
	
	for(int  i = 0; i < n; ++i){
		// counting the sum from 
		// 0 to i 
		sum += arr[i]; 
		
		// counting all the possible 
		// subarray which end at the 
		// index i 
		ans += (ll)mp[sum - (i + 1)]; 
		
		++mp[sum - (i + 1)]; 	
	}
	
	// returning the no of 
	// subarrays which length 
	// of the subarray 
	// is equal to sum of the 
	// subarray
	return ans; 
	
	
}
int main(){

         vector<int> arr = {0, 1, 1, 2}; 
		 
		cout << findSubarrays(arr) << endl; 
}

    

Output :

4