Need help with this question (Microsoft SDE Intern)

help
1

You are given array A of N integers and an integer S.Your task is to compute how many ways one can choose a sub-array that has arithmetic mean equal to S.

Constraints:
1<=N<=10^5
-10^9 <=S<=10^9
-10^9 <= A[i] <= 10^9

Examples:
A=[2,1,3]
S=2

Answer: 3
Explanation:
[2] ,[1,3] , [2,1,3]

Any help would be highly appreciated . Better than O(N*N) :sweat_smile:

2
1 Like
3

int solve (vector arr, int n, int avg) {
int res = 0, sum = 0;
unordered_map<int, int> mp;
mp[0] = 1;

for(int i = 0; i < n; i++){
    sum += arr[i] - avg;
    if(mp.find(sum) != mp.end()){
        res += mp[sum];
    }
    mp[sum]++;
}
return res;

}