LeetCode 930.和相同的二元子数组

Posted 2020-12-14 机器狗mo

在由若干 0 和 1 组成的数组 A 中，有多少个和为 S 的非空子数组。

示例：
输入：A = [1,0,1,0,1], S = 2
输出：4
解释：
如下面黑体所示，有 4 个满足题目要求的子数组：
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]

提示：
A.length <= 30000
0 <= S <= A.length
A[i] 为 0 或 1

class Solution:
    def numSubarraysWithSum(self, A: List[int], S: int) -> int:
        ans = 0
        pre = [0 for i in range(len(A)+1)]   
        pre_sum = 0
        pre[0] = 1    ##  前缀和为0的个数默认为1
        for i in range(len(A)):            
            pre_sum += A[i]        ## 更新前缀和
            print(i,pre,pre_sum - S,ans)
            if (pre_sum - S) >= 0:       ## 当前前缀和大于S，开始更新ans
                ans += pre[pre_sum - S]  ## 长前缀 减 短前缀 等于 子数组，
            pre[pre_sum] += 1       ## 记录前缀和为pre_sum的前缀个数，例如pre[3]=2 表示前缀和为3的前缀有2个
            print(i,pre,pre_sum - S,ans)
        return ans 

更为一般的情况：LeetCode 560.和为K的子数组（https://www.cnblogs.com/sandy-t/p/13227941.html）
题目如果有连续子数组求或乘积，就应该联想前缀数组