Binary Subarrays With Sum

medium prefix sum hash map sliding window

Problem

Given a binary array nums and an integer goal, return the number of non-empty subarrays with a sum equal to goal. A subarray is a contiguous part of the array.

Inputnums = [1, 0, 1, 0, 1], goal = 2

Output4

The four subarrays summing to 2 are [1,0,1], [1,0,1,0], [0,1,0,1] and [1,0,1] (indices 2..4).

nums (comma-separated 0/1)

goal

def num_subarrays_with_sum(nums, goal):
    count = {0: 1}
    prefix = 0
    total = 0
    for x in nums:
        prefix += x
        total += count.get(prefix - goal, 0)
        count[prefix] = count.get(prefix, 0) + 1
    return total

function numSubarraysWithSum(nums, goal) {
  const count = new Map([[0, 1]]);
  let prefix = 0, total = 0;
  for (const x of nums) {
    prefix += x;
    total += count.get(prefix - goal) || 0;
    count.set(prefix, (count.get(prefix) || 0) + 1);
  }
  return total;
}

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

int numSubarraysWithSum(vector<int>& nums, int goal) {
    unordered_map<int, int> count;
    count[0] = 1;
    int prefix = 0, total = 0;
    for (int x : nums) {
        prefix += x;
        if (count.count(prefix - goal)) total += count[prefix - goal];
        count[prefix]++;
    }
    return total;
}

Explanation

We want to count subarrays whose sum equals goal. The key idea is the running prefix sum: a subarray sums to goal exactly when two prefix sums differ by goal.

As we scan, prefix is the sum of everything so far. A subarray ending here has sum goal if some earlier prefix equalled prefix - goal. So we keep a hash map count that records how many times each prefix value has occurred.

At each element we add x to prefix, then add count[prefix - goal] to the answer (how many earlier start points give the right sum), and finally bump count[prefix]. We seed the map with {0: 1} so subarrays starting at index 0 are counted.

Example: nums = [1,0,1,0,1], goal = 2. The prefix sums are 1,1,2,2,3. Each time prefix - 2 has already been seen, we find a valid subarray; tallying these gives 4.

Because every lookup and update is constant time, the whole count is done in one pass.

Time: O(n) Space: O(n)