Description

Given a binary array nums and an integer goal, return the number of non-empty subarrays with a sum goal.

A subarray is a contiguous part of the array.

Example 1:

1Input: nums = [1,0,1,0,1], goal = 2
2Output: 4
3Explanation: The 4 subarrays are bolded and underlined below:
4[1,0,1]
5[1,0,1,0]
6[0,1,0,1]
7[1,0,1]

Example 2:

1Input: nums = [0,0,0,0,0], goal = 0
2Output: 15

Constraints:

  • 1 <= nums.length <= 3 * 10^4
  • nums[i] is either 0 or 1.
  • 0 <= goal <= nums.length

Solution

This problem could be solved using prefix sum operations. In this case, we want to find number of subarrays where sum of its elements is equal to goal. If we calculate the prefix sum for input array [1, 0, 1, 0, 1], we get [1, 1, 2, 2, 3].

1input: [1, 0, 1, 0, 1]
2prefixsum: [1, 1, 2, 2, 3]
3prepend 1: [1, 1, 1, 2, 2, 3]
4Even though we inserted 1 for key=0, currSum checks from index position 1
5Number of subarrays where sum will be 2 are below from prefix sum
6[0, 1, 1, 2]        => [1, 0, 1]
7[0, 1, 1, 2, 2]     => [1, 0, 1, 0]
8[1, 1, 2, 2, 3]     => [0, 1, 0, 1]
9[1, 2, 2, 3]        => [1, 0, 1]

 1class Solution {
 2    public int numSubarraysWithSum(int[] nums, int goal) {
 3        Map<Integer,Integer> map = new HashMap<>();
 4
 5        int currSum = 0, count = 0;
 6        map.put(0, 1);
 7        for (int i = 0; i < nums.length; i++) {
 8            currSum = currSum + nums[i];
 9            count += map.getOrDefault(currSum - goal,0);
10            map.put(currSum, map.getOrDefault(currSum, 0) + 1);
11            System.out.println(map + " count: " + count + " currSum: " + currSum);
12        }
13        return count;
14    }
15}
  • Time Complexity: O(n)
  • Space Complexity: O(n)