Solution


Approach 1: Index of Ones

Intuition

Say we number the 1s in A: with .

Then, if we have a subarray of sum , it has to use the ones . For each , we can count the number of such subarrays individually.

Algorithm

In general, the number of such subarrays (for ) is .

For example, if , then in A = [1,0,1,0,1,0,0,1], let's count the number of subarrays [i, j] that use the middle two 1s. There are 2 choices for the i (i = 1, 2) and 3 choices for the j (j = 4, 5, 6).

The corner cases are when , , or . We can handle these gracefully.

Complexity Analysis

  • Time Complexity: , where is the length of A.

  • Space Complexity: .


Approach 2: Prefix Sums

Intuition

Let P[i] = A[0] + A[1] + ... + A[i-1]. Then P[j+1] - P[i] = A[i] + A[i+1] + ... + A[j], the sum of the subarray [i, j].

Hence, we are looking for the number of i < j with P[j] - P[i] = S.

Algorithm

For each j, let's count the number of i with P[j] = P[i] + S. This is analogous to counting the number of subarrays ending in j with sum S.

It comes down to counting how many P[i] + S we've seen before. We can keep this count on the side to help us find the final answer.

Complexity Analysis

  • Time Complexity: , where is the length of A.

  • Space Complexity: .


Approach 3: Three Pointer

Intuition

For each j, let's try to count the number of i's that have the subarray [i, j] equal to S.

It is easy to see these i's form an interval [i_lo, i_hi], and each of i_lo, i_hi are increasing with respect to j. So we can use a "two pointer" style approach.

Algorithm

For each j (in increasing order), let's maintain 4 variables:

  • sum_lo : the sum of subarray [i_lo, j]
  • sum_hi : the sum of subarray [i_hi, j]
  • i_lo : the smallest i so that sum_lo <= S
  • i_hi : the largest i so that sum_hi <= S

Then, (provided that sum_lo == S), the number of subarrays ending in j is i_hi - i_lo + 1.

As an example, with A = [1,0,0,1,0,1] and S = 2, when j = 5, we want i_lo = 1 and i_hi = 3.

Complexity Analysis

  • Time Complexity: , where is the length of A.

  • Space Complexity: .


Analysis written by: @awice.