930. Binary Subarrays With Sum

https://leetcode.com/problems/binary-subarrays-with-sum/

In an array A of 0s and 1s, how many non-empty subarrays have sum S?

Example 1:

Input: A = [1,0,1,0,1], S = 2
Output: 4
Explanation: 
The 4 subarrays are bolded below:
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]

Note:

1. A.length <= 30000

2. 0 <= S <= A.length

3. A[i] is either 0 or 1

Thoughts

给定二进制数组有多少子数组和为S。子数组和问题 => presum或sliding window。presum没有什么特别的。sliding window不是直接求解满足条件的子数组和，而是先求和为[0:S]的所有子数组个数，再求和为[0:S-1]的所有子数组个数，它俩的差为子数组和为S的个数。用sliding window求和为[0:S]的所有子数组个数实现简单。

思路参考自这。

Subarray sum == S想到用hash map依次遍历把所有的 presum存起来，然后类似背包看count[presum - S]的个数，再+到res上

注意应先更新res, 因为是presum表之前有没有出现过这个sum，再更新count[presum]个数。

Code

class Solution:
    def numSubarraysWithSum(self, A: List[int], S: int) -> int:
        def at_most(S):
            l, s, res = 0, 0, 0
            for r, ar in enumerate(A):
                s += ar
                while s > S and l <= r:
                    s -= A[l]
                    l += 1
                res += r - l + 1
            return res
        return at_most(S) - at_most(S - 1)

class Solution {
public:
    int numSubarraysWithSum(vector<int>& A, int S) {
        unordered_map<int, int> count;
        int presum = 0, res = 0;
        count[0] = 1;
        for (int a : A) {
            presum += a;
            res += count[presum - S];
            ++count[presum];
        }
        return res;
    }
};