Count subarrays with equal number of 1’s and 0’s in C++

We are given an array arr[] containing 0’s and 1’s only. The goal is to count all subarrays of arr[] such that occurrences of 0s and 1s is equal in all. If the array is [1,0,0] .Subarray will be [1,0] only.

Let us understand with examples.

Input − arr[] = { 0, 0, 1, 1, 1, 0 };

Output − Count of subarrays with an equal number of 1’s and 0’s are − 4

Explanation − Subaarays will be −

arr[0 to 3] = [0,0,1,1],
arr[1 to 2] = [0,1],
arr[4 to 5] =[1,0],
Arr[0 to 5] =[0,0,1,1,1,0].

Input − arr[] = { 0, 1, 1, 1, 1 };

Output − Count of subarrays with equal number of 1’s and 0’s are − 1

Explanation − Subaarays will be − arr[0 to 1] = [0,1]

The approach used in the below program is as follows

We will traverse the array using two for loops to generate all subarrays possible. From i=0 to i<=size-1 and j=i to j<=size-1. Subarrays formed will be between arr[i] to arr[j]. Count frequency of 0 and 1 in each subarray. If equal then increment the count.

• Take an array arr[] of numbers.

• Function sub_zeroes_ones(int arr[], int size) takes the array and returns the count of subarrays with equal no. of 0’s and 1’s.

• Take the initial count as 0.

• We will traverse the array using two for loops from i=0 to i<=size-1 and j=0 to j<=size-1.

• Take two variables total_0, total_1 as 0 for the number of 0’s and 1’s in subarray arr[i] to arr[j].

• Compare arr[j] with 0 and 1. If arr[j] is 0 or 1 then increment respective count ( total_0 or total_1).

• If total_0==total_1. Increment count. ( the subarray has the same number of 0’s and 1’s as elements).

• At the end of both loops, return count as result.

Example

 Live Demo

#include <bits/stdc++.h>
using namespace std;
int sub_zeroes_ones(int arr[], int size){
   int count = 0;
   for (int i = 0; i <= size - 1; i++){
      int total_0 = 0;
      int total_1 = 0;
      for (int j = i; j <= size - 1; j++){
         if (arr[j] == 0){
            total_0++;
         }
         else if (arr[j] == 1){
            total_1++;
         }
         if(total_0 == total_1){
            count++;
         }
      }
   }
   return count;
}
int main(){
   int arr[] = {0, 1, 1, 0, 0};
   int size = sizeof(arr)/sizeof(arr[0]);
   cout<<"Count of subarrays with equal number of 1’s and 0’s are: "<<sub_zeroes_ones(arr, size);
}

Output

If we run the above code it will generate the following output −

Count of subarrays with equal number of 1’s and 0’s are: 4