Prefix sum problems: Given bianry array and a value k, return the count of subarray sum

Problem Statement:

Given a bianry array and a value k, you need to find the count of sub array sum whose sum is equal to k.

Example:

Input:  arr = [1, 0, 1, 0, 1] k = 2 
Output: 4

Explanation: 

[1, 0, 1]
[1, 0, 1, 0]
[-, 0, 1, 0, 1]
[-, -, 1, 0, 1]

Solution Explanation:

We will use hashmap along with prefix sum approach.

We will keep track of the sum while traversing the arrays.

Then store the frequency of the prefix sums, by subtracting the target sum from current sum.

Time Complexity: O(n)
Space Complexity: O(n)

Code Solution

#include <iostream>
#include <vector>
#include <climits>
#include <algorithm>
#include <unordered_map>
using namespace std;

int solution(vector<int> arr, int k) 
{
    int n = arr.size();
    int result = 0;
    int prefixSum = 0;

    unordered_map<int,int>mp;

    mp[0] = 1;

    for(int i=0; i<n; i++)
    {
        prefixSum += arr[i];
        result += mp[prefixSum - k];
        mp[prefixSum]++;
    }
    return result;
    
}

int main()
{
    vector<int> arr =  {1, 0, 1, 0, 1};
    int k = 2;
  
    cout<<solution(arr, k)<<endl;
    return 0;
}

Output

4