930

930. Binary Subarrays With Sum (Medium)

Given a binary array nums and an integer goal, return the number of non-empty subarrays with a sum goal.

A subarray is a contiguous part of the array.

 

Example 1:

Input: nums = [1,0,1,0,1], goal = 2
Output: 4
Explanation: The 4 subarrays are bolded and underlined below:
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]

Example 2:

Input: nums = [0,0,0,0,0], goal = 0
Output: 15

 

Constraints:

• 1 <= nums.length <= 3 * 104
• nums[i] is either 0 or 1.
• 0 <= goal <= nums.length

Companies:
C3 IoT, Adobe

Related Topics:
Array, Hash Table, Sliding Window, Prefix Sum

Solution 1. Prefix State Map

// OJ: https://leetcode.com/problems/binary-subarrays-with-sum/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(N)
class Solution {
public:
    int numSubarraysWithSum(vector<int>& A, int goal) {
        unordered_map<int, int> m{{0,-1}}; // count of 1s -> index
        int ans = 0;
        for (int i = 0, N = A.size(), sum = 0; i < N; ++i) {
            sum += A[i];
            if (m.count(sum) == 0) m[sum] = i;
            if (goal == 0) ans += i - m[sum];
            else if (sum - goal >= 0) ans += m[sum - goal + 1] - m[sum - goal];
        }
        return ans;
    }
};

Or

// OJ: https://leetcode.com/problems/binary-subarrays-with-sum/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(N)
class Solution {
public:
    int numSubarraysWithSum(vector<int>& A, int S) {
        unordered_map<int, int> m{{0, 1}}; // sum -> number of occurrences of this sum
        int sum = 0, ans = 0;
        for (int n : A) {
            sum += n;
            ans += m.count(sum - S) ? m[sum - S] : 0;
            m[sum]++;
        }
        return ans;
    }
};

Solution 2. At Most to Equal + Find Maximum Sliding Window

Check out "C++ Maximum Sliding Window Cheatsheet Template!".

// OJ: https://leetcode.com/problems/binary-subarrays-with-sum/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(1)
class Solution {
    int atMost(vector<int> A, int goal) {
        int N = A.size(), i = 0, j = 0, cnt = 0, ans = 0;
        while (j < N) {
            cnt += A[j++];
            while (i < j && cnt > goal) cnt -= A[i++];
            ans += j - i;
        }
        return ans;
    }
public:
    int numSubarraysWithSum(vector<int>& A, int goal) {
        return atMost(A, goal) - atMost(A, goal - 1);
    }
};

Similar problem:

• 992. Subarrays with K Different Integers (Hard)
• 1248. Count Number of Nice Subarrays (Medium)