Binary Subarrays With Sum Python Solution

lcpy gen -n 930   # by problem number
lcpy gen -s binary_subarrays_with_sum   # by problem name

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]

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

class Solution:
    # Time: O(n)
    # Space: O(n)
    def num_subarrays_with_sum(self, nums: list[int], goal: int) -> int:
        # Prefix sum counts: sum -> number of prefixes with that sum
        prefix_counts: dict[int, int] = {0: 1}
        total = 0
        count = 0
        for num in nums:
            total += num
            count += prefix_counts.get(total - goal, 0)
            prefix_counts[total] = prefix_counts.get(total, 0) + 1
        return count