Leetcode930 Binary Subarry With Sum

Posted on In

题目描述

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

Example 1:

1
2
3
4
5
6
7
8
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:

  • A.length <= 30000
  • 0 <= S <= A.length
  • A[i] is either 0 or 1

解析

题目是在一个由0、1组成的序列中，找到和为S的subarray的个数。

所有subarray的题目都可以考虑用prefix sum进行优化，即sum(i ,j) = prefix(j) - prefix(i)对于本题而言，可以利用HashMap存储prefix为特定sum的序列个数，每次有X新的元素加入的时候，更新即可，代码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
public int numSubarraysWithSum(int[] A, int S) {
    Map<Integer, Integer> map = new HashMap<>();
    map.put(0, 1);
    int sum = 0;
    int res = 0;
    
    for(int i = 0; i < A.length; i++){
        sum += A[i];
        int target = sum - S;
        //检查prefix和为target的序列有多少个
        res += map.getOrDefault(target, 0);
        map.put(sum, map.getOrDefault(sum ,0) + 1);
    }
    return res;
}