九章算法 | 腾讯面试题：和相同的二元子数组

描述

在由若干 0 和 1 组成的数组 A 中，有多少个和为 S 的非空子数组。

A.length <= 30000
0 <= S <= A.length
A[i] 为 0 或 1

[在线评测地址](https://www.lintcode.com/problem/binary-subarrays-with-sum/?utm_source=sc-zhihuzl-sy)

样例 1

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

样例 2

Input: A = [0,0,0,0,0,0,1,0,0,0], S = 0
Output: 27
Explanation: 
And 27 subarrays for S.

题解

前缀和：定义一个数组 sumN+1，si表示数组 A 中前 i 个元素之和，然后遍历 sum 数组，计算 si+S(含义：前 i 个元素之和是 si,找和为 S 的子数组个数)。求 si+S 的个数

public class Solution {
    /**
     * @param A: an array
     * @param S: the sum
     * @return: the number of non-empty subarrays
     */
    public int numSubarraysWithSum(int[] A, int S) {
        // Write your code here.
        if (A == null || A.length == 0) {
            return 0;
        }
        int prefixSum = 0;
        int res = 0;
        // counts[i] means the number of prefixSum = i
        int[] counts = new int[A.length + 1];
        counts[0] = 1;
        for (int i = 0; i < A.length; i++) {
            prefixSum += A[i];
            if (prefixSum >= S) {
                res += counts[prefixSum - S];
            }
            counts[prefixSum]++;
        }
        return res;
    }
}