Sliding Window: Return count of subarray with sum equal to k

Problem Statement:

You are given a binary array and an integer k.

You need to return the count of subarray with sum equal to k.

Example:

Input: arr[] = [1, 0, 1, 1, 0 ], K = 2

Output: 4

Explanation: 
valid subarray = 4
[1, 0, 1]
[1, 1]
[0, 1, 1]
[1, 1, 0]
[0,1, 1, 0]

Solution Explanation:

We will use sliding window in this approach with 2 pointers.

We have a main function, that will call the helper function.

We call the helper function with subarray sum equal to goal and goal -1.

This will give the solution for us.

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

Code Solution

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

int atMost(vector<int>& nums, int goal)
{
    int window_start=0, window_end=0, ans=0,count=0;

    if (goal < 0)
        return 0;

    for(window_end=0; window_end<nums.size() ;window_end++)
    {
        count+=nums[window_end];
        while(count>goal)
        {
            count-=nums[window_start];
            window_start++;
        }
        ans+=window_end-window_start+1;
    }
    return ans;
}

int solution(vector<int>& nums, int goal) 
{
    return atMost(nums,goal) - atMost(nums,goal-1);
}

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

    cout << solution(arr, k) << endl;

    return 0;
}

Output

5