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

在由若干 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)
题目如果有连续子数组求或乘积,就应该联想前缀数组

posted @ 2020-07-02 23:28  机器狗mo  阅读(193)  评论(0编辑  收藏  举报