Formatted question description: https://leetcode.ca/all/930.html
930. Binary Subarrays With Sum
Level
Medium
Description
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:
A.length <= 300000 <= S <= A.lengthA[i]is either0or1.
Solution
Use a map to count the number of subarrays starting from index 0 that have sum equal to a given value.
If S is 0, then for each possible sum in the map, count the number of subarrays of the sum. If there are k such subarrays, then the number of non-empty subarrays obtained from such subarrays is k * (k - 1) / 2. Count the number of non-empty subarrays with sum 0 from all possible sums in the map.
If S is positive, then for each pair (i, i + S) such that i >= 0 and i + S doesn’t exceed the maximum possible sum (which is the sum of all numbers in the array), let count1 and count2 be the numbers of subarrays starting from index 0 that have sum i and sum i + S respectively, and the number of non-empty subarrays with sum S for the pair (i, i + S) is count1 * count2. Count the number of non-empty subarrays with sum S from all possible sums in the map.
class Solution {
public int numSubarraysWithSum(int[] A, int S) {
if (A == null || A.length == 0)
return 0;
Map<Integer, Integer> sumCountMap = new HashMap<Integer, Integer>();
int length = A.length;
int beginIndex = 0;
int curSum = A[0];
for (int i = 1; i < length; i++) {
if (A[i] == 1) {
sumCountMap.put(curSum, i - beginIndex);
curSum++;
beginIndex = i;
}
}
sumCountMap.put(curSum, length - beginIndex);
int totalCount = 0;
sumCountMap.put(0, sumCountMap.getOrDefault(0, 0) + 1);
if (S == 0) {
Set<Integer> keySet = sumCountMap.keySet();
for (int num : keySet) {
int count = sumCountMap.get(num);
totalCount += count * (count - 1) / 2;
}
return totalCount;
} else {
int minSum = A[0], maxSum = curSum;
int upper = maxSum - S;
for (int i = 0; i <= upper; i++) {
int count1 = sumCountMap.get(i), count2 = sumCountMap.get(i + S);
totalCount += count1 * count2;
}
return totalCount;
}
}
}