Problem Description

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

A subarray is defined as a contiguous sequence of elements within the array. For example, if nums = [1, 0, 1, 0, 1], then [1, 0, 1], [0, 1, 0], and [1] are all valid subarrays, but [1, 1] (skipping the 0 in between) is not.

The solution uses a prefix sum approach with a hash map. As we traverse the array, we maintain a running sum s of all elements seen so far. For each position, we check how many previous prefix sums equal s - goal. If a previous prefix sum equals s - goal, it means the subarray between that previous position and the current position sums to goal.

The hash map cnt stores the frequency of each prefix sum encountered. We initialize it with {0: 1} to handle cases where a prefix sum itself equals the goal. For each element in the array:

1. Add the current element to the running sum s
2. Check how many times we've seen the prefix sum s - goal before (this gives us the count of subarrays ending at the current position with sum equal to goal)
3. Update the hash map with the current prefix sum

The algorithm efficiently counts all valid subarrays in a single pass through the array with O(n) time complexity.

Intuition

The key insight comes from recognizing that we need to find all subarrays with a specific sum. A naive approach would check every possible subarray, but that would be inefficient with O(n²) or O(n³) complexity.

Let's think about what makes a subarray sum to our goal. If we know the sum from the start of the array to position i (let's call it prefix_sum[i]) and the sum from the start to position j where j < i (call it prefix_sum[j]), then the sum of the subarray from position j+1 to i equals prefix_sum[i] - prefix_sum[j].

For this subarray to have sum equal to goal, we need: prefix_sum[i] - prefix_sum[j] = goal

Rearranging this equation: prefix_sum[j] = prefix_sum[i] - goal

This means at position i, we need to find how many previous positions j have a prefix sum equal to prefix_sum[i] - goal. Each such position gives us one valid subarray.

Why use a hash map? As we traverse the array, we need to quickly look up how many times we've seen a specific prefix sum value before. A hash map provides O(1) lookup time for this count.

The initialization {0: 1} handles the edge case where a prefix sum itself equals the goal. For instance, if the first few elements sum to goal, we need to count this as one valid subarray. The "virtual" prefix sum of 0 before the array starts enables this counting.

This transforms our problem from checking all subarrays to a single pass through the array while maintaining prefix sum frequencies - a classic pattern for subarray sum problems.

Learn more about Prefix Sum and Sliding Window patterns.

Solution Approach

The solution implements a prefix sum technique with a hash map to efficiently count subarrays with the target sum.

Data Structures Used:

• Counter (hash map): Stores the frequency of each prefix sum encountered
• Two variables: ans for the result count and s for the running prefix sum

Implementation Steps:

1. Initialize the hash map: cnt = Counter({0: 1})

   • We start with {0: 1} to handle cases where a prefix sum itself equals goal
   • This represents the "empty prefix" before the array starts

2. Initialize tracking variables: ans = s = 0

   • ans: Counts the total number of valid subarrays
   • s: Maintains the running sum (prefix sum) as we traverse

3. Traverse through each element: for v in nums:

   • For each element v in the array:

4. Update the prefix sum: s += v

   • Add the current element to our running sum
   • Now s represents the sum from the start of the array to the current position

5. Count valid subarrays ending at current position: ans += cnt[s - goal]

   • Look up how many times we've seen the prefix sum s - goal
   • Each occurrence represents a valid subarray ending at the current position
   • If s - goal doesn't exist in cnt, it returns 0 (default behavior)

6. Update the frequency map: cnt[s] += 1

   • Record that we've seen the current prefix sum s one more time
   • This will be used for counting subarrays in future iterations

7. Return the result: return ans

Example Walkthrough: For nums = [1, 0, 1, 0, 1] and goal = 2:

• Start: cnt = {0: 1}, s = 0, ans = 0
• Index 0 (v=1): s = 1, check cnt[1-2] = cnt[-1] = 0, ans = 0, cnt = {0: 1, 1: 1}
• Index 1 (v=0): s = 1, check cnt[1-2] = cnt[-1] = 0, ans = 0, cnt = {0: 1, 1: 2}
• Index 2 (v=1): s = 2, check cnt[2-2] = cnt[0] = 1, ans = 1, cnt = {0: 1, 1: 2, 2: 1}
• Index 3 (v=0): s = 2, check cnt[2-2] = cnt[0] = 1, ans = 2, cnt = {0: 1, 1: 2, 2: 2}
• Index 4 (v=1): s = 3, check cnt[3-2] = cnt[1] = 2, ans = 4, cnt = {0: 1, 1: 2, 2: 2, 3: 1}

The algorithm runs in O(n) time complexity with O(n) space complexity for the hash map.

Example Walkthrough

Let's walk through a small example with nums = [1, 0, 1, 1] and goal = 2.

Initial Setup:

• cnt = {0: 1} (hash map with initial prefix sum 0)
• s = 0 (running sum)
• ans = 0 (count of valid subarrays)

Step 1: Process nums[0] = 1

• Update running sum: s = 0 + 1 = 1
• Check for valid subarrays: Look up cnt[s - goal] = cnt[1 - 2] = cnt[-1] = 0
  • No previous prefix sum equals -1, so no valid subarrays end here
• Update answer: ans = 0 + 0 = 0
• Record current prefix sum: cnt[1] = 1
• State: cnt = {0: 1, 1: 1}, ans = 0

Step 2: Process nums[1] = 0

• Update running sum: s = 1 + 0 = 1
• Check for valid subarrays: cnt[1 - 2] = cnt[-1] = 0
  • Still no valid subarrays
• Update answer: ans = 0 + 0 = 0
• Record current prefix sum: cnt[1] = 2 (we've seen prefix sum 1 twice now)
• State: cnt = {0: 1, 1: 2}, ans = 0

Step 3: Process nums[2] = 1

• Update running sum: s = 1 + 1 = 2
• Check for valid subarrays: cnt[2 - 2] = cnt[0] = 1
  • Found 1 valid subarray! This means the subarray from start to current position sums to 2
  • The subarray is [1, 0, 1]
• Update answer: ans = 0 + 1 = 1
• Record current prefix sum: cnt[2] = 1
• State: cnt = {0: 1, 1: 2, 2: 1}, ans = 1

Step 4: Process nums[3] = 1

• Update running sum: s = 2 + 1 = 3
• Check for valid subarrays: cnt[3 - 2] = cnt[1] = 2
  • Found 2 valid subarrays! Each previous occurrence of prefix sum 1 gives us a valid subarray
  • These correspond to subarrays [0, 1, 1] (from index 1 to 3) and [1, 1] (from index 2 to 3)
• Update answer: ans = 1 + 2 = 3
• Record current prefix sum: cnt[3] = 1
• Final state: cnt = {0: 1, 1: 2, 2: 1, 3: 1}, ans = 3

Result: The algorithm found 3 subarrays that sum to 2:

1. [1, 0, 1] (indices 0-2)
2. [0, 1, 1] (indices 1-3)
3. [1, 1] (indices 2-3)

The key insight is that when we're at position i with prefix sum s, finding cnt[s - goal] tells us exactly how many subarrays ending at position i have sum equal to goal.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n) where n is the length of the input array nums.

The algorithm iterates through the array exactly once with a single for loop. Within each iteration, it performs:

• Addition operation: s += v - O(1)
• Dictionary lookup: cnt[s - goal] - O(1) average case
• Dictionary update: cnt[s] += 1 - O(1) average case
• Addition to answer: ans += ... - O(1)

Since all operations inside the loop are constant time, and the loop runs n times, the overall time complexity is O(n).

Space Complexity: O(n) in the worst case.

The space is primarily used by the Counter dictionary cnt. In the worst case scenario where all prefix sums are unique (for example, when the input array contains all 1s and goal is larger than n), the dictionary will store n + 1 key-value pairs (including the initial {0: 1} and one entry for each prefix sum). Therefore, the space complexity is O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Forgetting to Initialize the Hash Map with {0: 1}

The Pitfall: A common mistake is initializing an empty hash map without the base case {0: 1}. This causes the algorithm to miss subarrays that start from index 0 and sum to the goal.

Why It Happens: When a prefix sum equals the goal (meaning current_sum - goal = 0), we need to count this as a valid subarray starting from index 0. Without {0: 1} in the initial hash map, these subarrays won't be counted.

Incorrect Code:

def numSubarraysWithSum(self, nums: List[int], goal: int) -> int:
    prefix_sum_count = Counter()  # Wrong! Missing {0: 1}
    result = 0
    current_sum = 0
  
    for num in nums:
        current_sum += num
        result += prefix_sum_count[current_sum - goal]
        prefix_sum_count[current_sum] += 1
  
    return result

Example Where It Fails:

• Input: nums = [1, 0, 1], goal = 1
• Without {0: 1}: Returns 2 (missing the subarray [1] at index 0)
• With {0: 1}: Returns 4 (correctly counts all subarrays)

Correct Solution:

prefix_sum_count = Counter({0: 1})  # Correct initialization

2. Updating the Hash Map Before Checking

The Pitfall: Another mistake is updating prefix_sum_count[current_sum] before checking prefix_sum_count[current_sum - goal]. This can lead to incorrect counting when goal = 0.

Incorrect Code:

for num in nums:
    current_sum += num
    prefix_sum_count[current_sum] += 1  # Wrong order!
    result += prefix_sum_count[current_sum - goal]

Why It's Wrong: When goal = 0, we have current_sum - goal = current_sum. If we increment the count first, we might count the current position itself as a valid "previous" position, which would be incorrect for a single element.

Example Where It Fails:

• Input: nums = [0, 0, 0], goal = 0
• With wrong order: Overcounts subarrays
• With correct order: Properly counts each valid subarray

Correct Solution: Always check first, then update:

result += prefix_sum_count[current_sum - goal]  # Check first
prefix_sum_count[current_sum] += 1              # Then update