[leetcode] 930. Binary Subarrays With Sum

Description

In an array A of 0s and 1s, how many non-empty subarrays have sum S?

Example 1:

Input: A = [1,0,1,0,1], S = 2
Output: 4
Explanation: 
The 4 subarrays are bolded below:
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]

Note:

1. A.length <= 30000
2. 0 <= S <= A.length
3. A[i] is either 0 or 1.

analysis

The meaning of the question is: given an array, find the sum as a sub-array of a given S. For this question, I refer to the idea of ​​the answer, using the preSum method:

• First use preSum to sum each position of the array, and preSum=A[0] + A[1] +… + A[i-1],
• Then preSum[j+1]-preSum[i] = A[i] + A[i+1] +… + A[j], which is the sum of the intervals of the array [i, j].
• Then we record the number of preSum[j] = preSum[i] + S is the final answer, which is equivalent to finding the number ending in j and the sum is S.

For example: A = [1,0,1,0,1], S = 2

defaultdict(<class 'int'>, {0: 0, 2: 1, 1: 0, 3: 2, 4: 2, 5: 1}) 
其中key是preSum的值，value为频率

It's a bit difficult to understand, you need to do a manual calculation.
I started to use the double pointer to solve the zero situation.

Code

class Solution:
    def numSubarraysWithSum(self, A: List[int], S: int) -> int:
        res=0
        preSum=[0]
        for x in A:
            preSum.append(preSum[-1]+x)
        d=defaultdict(int)
        for x in preSum:
            res+=d[x]
            d[x+S]+=1
            
        return res

references

Approach 2: Prefix Sums