動機

用atMost秒殺 用prefix sum in fly也可以

Problem

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

\r\r

A subarray is a contiguous part of the array.

\r\r

 

\r

Example 1:

\r\r
\rInput: nums = [1,0,1,0,1], goal = 2\rOutput: 4\rExplanation: The 4 subarrays are bolded and underlined below:\r[1,0,1,0,1]\r[1,0,1,0,1]\r[1,0,1,0,1]\r[1,0,1,0,1]\r
\r\r

Example 2:

\r\r
\rInput: nums = [0,0,0,0,0], goal = 0\rOutput: 15\r
\r\r

 

\r

Constraints:

\r\r
    \r
  • 1 <= nums.length <= 3 * 104
    • \r
  • nums[i] is either 0 or 1.
    • \r
  • 0 <= goal <= nums.length
    • \r

Sol

class Solution:
    def numSubarraysWithSum(self, nums: List[int], goal: int) -> int:
        def atmost(bound):
            ret = cnts = i = 0
            for j,n in enumerate(nums):
                cnts += n
                while i <= j and cnts > bound:
                    cnts -= nums[i]
                    i += 1
                ret += j-i+1
            return ret
        return atmost(goal)-atmost(goal-1)

def numSubarraysWithSum(self, A, S):
    c = collections.Counter({0: 1})
    psum = res = 0
    for i in A:
        psum += i
        res += c[psum - S]
        c[psum] += 1
    return res