930

930. Binary Subarrays With Sum

We’ve got an array of 1s and 0s, and we need to find all the subarrays that sum to the goal provided. One would think a sliding window approach would work, but actually we want to use a prefix sum. We keep track of the current sum of the numbers we’ve seen, and index into a dictionary that keeps track of the number of times we’ve seen a particular sum. We add to the answer the number of times we’ve seen the current sum minus the goal - providing us with a number of subarrays that match that constraint.

The solution is as follows:

from collections import defaultdict
 
class Solution:
    def numSubarraysWithSum(self, nums: List[int], goal: int) -> int:
        prefix = defaultdict(int)
        prefix[0] = 1
        ans = curr = 0
 
        for num in nums:
            curr += num
            ans += prefix[curr - goal]
            prefix[curr] += 1
 
        return ans

_ Time Complexity:

O(n) - We iterate through the input once.

_ Space Complexity:

O(n) - We maintain a dictionary of prefix sums.