LeetCode-python 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

解题思路

---

思路很像LeetCode-python 560.和为K的子数组
由于只有0和1，可以用数组代替字典
pre_sum[i]表示和为i的前缀的个数

代码实现

class Solution(object):
    def numSubarraysWithSum(self, A, S):
        """
        :type A: List[int]
        :type S: int
        :rtype: int
        """
        pre_sum = [0]*(len(A)+1)
        pre_sum[0] = 1
        cur_sum = 0
        res = 0
        for num in A:
            cur_sum += num
            if pre_sum[cur_sum - S]:
                res += pre_sum[cur_sum - S]
            pre_sum[cur_sum] += 1
        return res

本文链接：https://www.jianshu.com/p/2efa08876465