Problem

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

A subarray is a contiguous part of the array.

Examples

Example 1

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

Example 2

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

Constraints

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

Solution

Method 1 – Prefix Sum and Hash Map

Intuition

We want to count the number of subarrays whose sum is exactly goal. By using a prefix sum and a hash map to count the number of times each prefix sum occurs, we can efficiently count the number of valid subarrays ending at each index.

Approach

  1. Initialize a hash map to store the count of each prefix sum, starting with 0 mapped to 1.
  2. Iterate through the array, maintaining a running sum.
  3. For each element, add the count of (current sum - goal) to the answer.
  4. Update the hash map with the current sum.
  5. Return the answer.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
class Solution {
public:
    int numSubarraysWithSum(vector<int>& nums, int goal) {
        unordered_map<int, int> cnt;
        cnt[0] = 1;
        int sum = 0, ans = 0;
        for (int x : nums) {
            sum += x;
            ans += cnt[sum - goal];
            cnt[sum]++;
        }
        return ans;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
func numSubarraysWithSum(nums []int, goal int) int {
    cnt := map[int]int{0: 1}
    sum, ans := 0, 0
    for _, x := range nums {
        sum += x
        ans += cnt[sum-goal]
        cnt[sum]++
    }
    return ans
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Solution {
    public int numSubarraysWithSum(int[] nums, int goal) {
        Map<Integer, Integer> cnt = new HashMap<>();
        cnt.put(0, 1);
        int sum = 0, ans = 0;
        for (int x : nums) {
            sum += x;
            ans += cnt.getOrDefault(sum - goal, 0);
            cnt.put(sum, cnt.getOrDefault(sum, 0) + 1);
        }
        return ans;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Solution {
    fun numSubarraysWithSum(nums: IntArray, goal: Int): Int {
        val cnt = mutableMapOf(0 to 1)
        var sum = 0
        var ans = 0
        for (x in nums) {
            sum += x
            ans += cnt.getOrDefault(sum - goal, 0)
            cnt[sum] = cnt.getOrDefault(sum, 0) + 1
        }
        return ans
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
class Solution:
    def numSubarraysWithSum(self, nums: list[int], goal: int) -> int:
        from collections import defaultdict
        cnt: dict[int, int] = defaultdict(int)
        cnt[0] = 1
        s = ans = 0
        for x in nums:
            s += x
            ans += cnt[s - goal]
            cnt[s] += 1
        return ans
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
impl Solution {
    pub fn num_subarrays_with_sum(nums: Vec<i32>, goal: i32) -> i32 {
        use std::collections::HashMap;
        let mut cnt = HashMap::new();
        cnt.insert(0, 1);
        let (mut sum, mut ans) = (0, 0);
        for x in nums {
            sum += x;
            ans += *cnt.get(&(sum - goal)).unwrap_or(&0);
            *cnt.entry(sum).or_insert(0) += 1;
        }
        ans
    }
}

Complexity

  • ⏰ Time complexity: O(n) — Each element is processed once.
  • 🧺 Space complexity: O(n) — For the hash map storing prefix sums.