leetcode [930] 和相同的二元子数组

---

Contact me:

Blog -> https://cugtyt.github.io/blog/index
Email -> cugtyt@qq.com
GitHub -> Cugtyt@GitHub

---

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

来自题解:

我们遍历区间的右端点 j，同时维护四个变量：

sum_lo：A[i_lo..j] 的值；

sum_hi：A[i_hi..j] 的值；

i_lo：最小的满足 sum_lo <= S 的 i；

i_hi：最大的满足 sum_hi <= S 的 i。

例如，当数组 A 为 [1,0,0,1,0,1]，S 的值为 2 。当 j = 5 时，i_lo 的值为 1，i_hi 的值为 3。对于每一个 j，和为 S 的非空子数组的数目为 i_hi - i_lo + 1。

class Solution:
    def numSubarraysWithSum(self, A: List[int], S: int) -> int:
        i_lo = i_hi = 0
        sum_lo = sum_hi = 0
        ans = 0
        for j, x in enumerate(A):
            # Maintain i_lo, sum_lo:
            # While the sum is too big, i_lo += 1
            sum_lo += x
            while i_lo < j and sum_lo > S:
                sum_lo -= A[i_lo]
                i_lo += 1

            # Maintain i_hi, sum_hi:
            # While the sum is too big, or equal and we can move, i_hi += 1
            sum_hi += x
            while i_hi < j and (
                    sum_hi > S or sum_hi == S and A[i_hi] == 0):
                sum_hi -= A[i_hi]
                i_hi += 1

            if sum_lo == S:
                ans += i_hi - i_lo + 1

        return ans