930. Binary Subarrays With Sum (Medium) - TengnanYao/daily_leetcode GitHub Wiki

class Solution(object):
    def numSubarraysWithSum(self, nums, goal):
        """
        :type nums: List[int]
        :type goal: int
        :rtype: int
        """
        # prefix sums
        h = collections.Counter({0: 1})
        psum = result = 0
        for num in nums:
            psum += num
            result += h[psum - goal]
            h[psum] += 1
        return result

        # combinations of 0
        result = 0
        ones = [-1] + [i for i in range(len(nums)) if nums[i]] + [len(nums)]
        if goal == 0:
            for i in range(len(ones) - 1):
                m = ones[i + 1] - ones[i] - 1
                result += m * (m + 1) // 2
        else:
            for i in range(len(ones) - 1 - goal):
                result += (ones[i + 1] - ones[i]) * (ones[i + goal + 1] - ones[i + goal])
        return result