Code Implementation

class Solution:
    def numSubarraysWithSum(self, nums, goal):
        count = 0
        prefix_counts = {0: 1}
        prefix_sum = 0
        for num in nums:
            prefix_sum += num
            count += prefix_counts.get(prefix_sum - goal, 0)
            prefix_counts[prefix_sum] = prefix_counts.get(prefix_sum, 0) + 1
        return count
      

class Solution {
public:
    int numSubarraysWithSum(vector<int>& nums, int goal) {
        unordered_map<int, int> prefix_counts;
        prefix_counts[0] = 1;
        int prefix_sum = 0, count = 0;
        for (int num : nums) {
            prefix_sum += num;
            if (prefix_counts.find(prefix_sum - goal) != prefix_counts.end()) {
                count += prefix_counts[prefix_sum - goal];
            }
            prefix_counts[prefix_sum]++;
        }
        return count;
    }
};
      

class Solution {
    public int numSubarraysWithSum(int[] nums, int goal) {
        Map<Integer, Integer> prefixCounts = new HashMap<>();
        prefixCounts.put(0, 1);
        int prefixSum = 0, count = 0;
        for (int num : nums) {
            prefixSum += num;
            count += prefixCounts.getOrDefault(prefixSum - goal, 0);
            prefixCounts.put(prefixSum, prefixCounts.getOrDefault(prefixSum, 0) + 1);
        }
        return count;
    }
}
      

var numSubarraysWithSum = function(nums, goal) {
    let prefixCounts = {0: 1};
    let prefixSum = 0, count = 0;
    for (let num of nums) {
        prefixSum += num;
        count += (prefixCounts[prefixSum - goal] || 0);
        prefixCounts[prefixSum] = (prefixCounts[prefixSum] || 0) + 1;
    }
    return count;
};
      

Problem Description

You are given a binary array nums (an array containing only 0s and 1s) and an integer goal. Your task is to count the number of non-empty, contiguous subarrays of nums whose sum is exactly equal to goal.

• Each subarray must be contiguous (elements are consecutive in the original array).
• Subarrays can overlap but must not be empty.
• All elements are either 0 or 1.
• Return the total number of such subarrays.

For example, if nums = [1,0,1,0,1] and goal = 2, the answer would be 4.

Thought Process

The most straightforward way to solve this problem is to consider all possible subarrays, calculate their sums, and count those whose sum equals goal. This brute-force approach, however, is inefficient for large arrays, as it checks every possible start and end index, resulting in a time complexity of O(n²).

To optimize, we look for patterns or properties we can exploit. Since the array contains only 0s and 1s, the sum of any subarray is simply the count of 1s in that subarray. This property makes it possible to use prefix sums and hash maps to efficiently count the number of subarrays with a given sum.

The key insight is: for each position in the array, if we know how many times a prefix sum of current_sum - goal has occurred before, we can add that many subarrays ending at the current position with sum equal to goal.

Solution Approach

We use a prefix sum approach combined with a hash map to efficiently count the number of subarrays whose sum equals goal. Here’s how it works step by step:

1. Initialize:
   • A hash map (dictionary) prefix_counts to keep track of how many times each prefix sum has occurred, starting with 0: 1 to account for subarrays starting at index 0.
   • A variable prefix_sum to store the running sum as we iterate through nums.
   • A counter count to store the answer.
2. Iterate through the array:
   • For each number in nums, add it to prefix_sum.
   • Check if prefix_sum - goal exists in prefix_counts. If it does, add the number of times it has occurred to count. This step finds all subarrays ending at the current index with sum equal to goal.
   • Increment the count of prefix_sum in prefix_counts.
3. Return count as the answer.

Why use a hash map? The hash map allows us to check in constant time how many times a particular prefix sum has occurred, which is crucial for efficiency.

Example Walkthrough

Let's walk through the example nums = [1,0,1,0,1] with goal = 2:

1. Initialize: prefix_counts = {0:1}, prefix_sum = 0, count = 0
2. Index 0, num = 1: prefix_sum = 1. prefix_sum - goal = -1 (not in prefix_counts). Update prefix_counts to {0:1, 1:1}.
3. Index 1, num = 0: prefix_sum = 1. prefix_sum - goal = -1 (not in prefix_counts). Update prefix_counts to {0:1, 1:2}.
4. Index 2, num = 1: prefix_sum = 2. prefix_sum - goal = 0 (in prefix_counts, count is 1). Add 1 to count (count = 1). Update prefix_counts to {0:1, 1:2, 2:1}.
5. Index 3, num = 0: prefix_sum = 2. prefix_sum - goal = 0 (in prefix_counts, count is 1). Add 1 to count (count = 2). Update prefix_counts to {0:1, 1:2, 2:2}.
6. Index 4, num = 1: prefix_sum = 3. prefix_sum - goal = 1 (in prefix_counts, count is 2). Add 2 to count (count = 4). Update prefix_counts to {0:1, 1:2, 2:2, 3:1}.

Final answer: count = 4. The subarrays are [1,0,1], [0,1,0,1], [1,0,1,0], and [0,1].

Time and Space Complexity

• Brute-force approach:
  • Time Complexity: O(n²), since we check all possible subarrays.
  • Space Complexity: O(1), only a few variables are used.
• Optimized prefix sum + hash map approach:
  • Time Complexity: O(n), since we traverse the array once and all hash map operations are constant time.
  • Space Complexity: O(n), for storing prefix sums in the hash map (in the worst case, all prefix sums are unique).

Summary

The problem of counting binary subarrays with a given sum can be solved efficiently by leveraging prefix sums and a hash map. The key insight is that the number of subarrays ending at each position with sum equal to goal is determined by how many times the prefix sum current_sum - goal has occurred so far. This approach reduces the time complexity from quadratic to linear, making it suitable for large inputs.