How to count the subarrays with an equal number of 0s and 1s

Given a fixed-size array containing $N$ elements, all of which are either 1 or 0, we can find the number of subarrays in it that contain an equal number of 0s and 1s. For example, the array $[1, 0, 0, 1, 0, 1, 1]$ has a total of 8 subarrays with an equal number of 0s and 1s.

To do so, one possible solution is to generate all possible subarrays and count the number of 0s and 1s in each. However, this is a very inefficient solution as it will take $O(N^2)$ time in the worst case (this is because generating all possible subarrays requires one loop to be nested inside the other). In this Answer, we'll explore how a more efficient solution can be implemented for this task, which takes $O(N)$time in the worst case.

Solution

• To simplify this problem of returning the count equal_subarrays_count of the subarrays with an equal number of 0s and 1s, each time we encounter a 0 in the array while traversing it, we can consider it to be a -1.

  • By doing this, we can now keep track of the sum of all the elements in the array which we have traversed so far in a variable current_sum, and we also keep track of the number of times each value of current_sum is encountered.

• Each time we encounter a 1, we add 1 to current_sum, and each time we encounter a 0, we subtract 1 from current_sum (since we're considering 0 as -1).

• Now, if we ever encounter a value of current_sum, which has already been encountered before (it's count is greater than 0), this means that between the indices of the array at which this same value of current_sum occurs, the elements visited form a sub-array that has an equal number of 0s and 1s.

  • This is because each element encountered can only increment or decrement current_sum by 1, and hence, if current_sum is the same at any 2 indices, it means that current_sum was incremented the same number of times it was decremented. This can only happen if the same number of 0s and 1s were encountered between those indices.

• Since the count of each value of current_sum determines the number of indices between which the elements sum to the same current_sum, we then add this count of current_sum to equal_subarrays_count (since this count determines the number of new subarrays that can be formed containing an equal number of 0s and 1s).

• We then increment this count by 1 to mark another instance of it occurring.

This process is repeated until we have traversed the whole array, after which we can simply return equal_subarrays_count. The example below demonstrates how this solution works:

Since this solution only requires us to traverse the entire array once, it takes $O(N)$ time in the worst case.

Example

The code below shows how the algorithm discussed above can be implemented in Python:

Explanation

The countEqualSubarrays() function is explained below:

• Line 4: We store the number of times each value of current_sum occurs in a dictionary, with the key being the value of current_sum and the value is its count.

• Lines 9–12: We increment current_sum if we encounter a 1, or decrement it if we encounter a 0.

• Lines 14–16: We add the count of current_sum to equal_subarrays_count and then increment this count by 1.

• Lines 17–18: If the value of current_sum is seen for the first time, we make an entry for it in the dictionary and assign its count the value 1.