难度：★★★☆☆
类型：数组
方法：前缀和

题目

力扣链接请移步本题传送门
更多力扣中等题的解决方案请移步力扣中等题目录

在由若干 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

解答

我们用前缀和来解决这个问题。

定义一个数组P，比输入A多一个元素，P[i]表示A[:i]数组中所有元素之和。

前缀和的好处是，可以快速求取任意一段连续子数组之和，例如从sum(A[i:j])=P[j]-P[i]。

根据这个原理，我们就可以统计满足条件的子数组的个数了。最容易想到的办法是两层循环，遍历i和j来统计P[j]-P[i]==S的情况的个数，但是时间复杂度很高，我们用另一种做法。

准备一个计数器count，这是一个字典，字典的键代表的是j位置处的元素，字典的值表示以j位置处元素结尾的并且和为S的子数组的个数。

接下来是两个问题，一个是如何获得count的问题，一个是如何使用count的问题。

如何获得count，很容易想到，对于j位置处结尾的元素，我们可以通过遍历小于j的所有位置i，及时探测是否存在P[i]+S=P[j]的情况存在，每次等式成立，都可以将计数器中A[j]所对应的值加一。

如何使用count，同样需要遍历原数组A，对于其中的每个元素x，一旦在count中存在，则将其在计数器中对应的数值加入到结果中，这样做的含义是，清算以x结尾的满足条件的子数组的个数。

上述两个过程可以在同一个循环中实现，但是为了应对S为零的情况，需要先使用count后更新count，两个过程不冲突。

from collections import Counter


class Solution:
    def numSubarraysWithSum(self, A, S):
        P = [0]
        for x in A:
            P.append(P[-1] + x)
        count = Counter()

        ans = 0
        for x in P:
            ans += count[x]
            count[x + S] += 1

        return ans

如有疑问或建议，欢迎评论区留言~

有关更多力扣中等题的python解决方案，请移步力扣中等题解析