Count subarrays consisting of only 0’s and only 1’s in a binary array 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 each subarray contains only 0’s or only 1’s not both. If the array is [1,0,0] .Subarrays will be for 0’s only [0], [0], [0,0] and for 1’s only [1].

Let us understand with examples.

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

Output − Subarrays with only 0's are : 4 Subarrays with only 1's are : 6

Explanation − Subaarays will be −

For all 0’s: [0], [0], [0], [0,0]. Total 4 ( arr[0], arr[1], arr[5], arr[0-1] )
For all 1’s: [1], [1], [1], [1,1], [1,1], [1,1,1]. Total 6 ( arr[2], arr[2], arr[4], arr[2-3], arr[3-4],
arr[2-4] )

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

Output − Subarrays with only 0's are : 2 Subarrays with only 1's are : 2

Explanation − Subaarays will be −

For all 0’s: [0], [0]. Total 2 ( arr[1], arr[3] )
For all 1’s: [1], [1]. Total 2 ( arr[0], arr[2] )

The approach used in the below program is as follows

We will traverse the array twice for separately counting subarrays with 0’s and 1’s only. Take two counters count_0 and count_1 to store the count of consecutive 0’s and 1’s in the array. For each such count, possible subarrays would be count_0*(count_0+1)/2 and similarly for count_1.

Add this to total_0 count until the end of the array is reached. Do a similar process for 1’s.

• Take an array arr[] of numbers.

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

• Take the initial counts as temp_0 and temp_1 for subarrays.

• Take the temporary consecutive counts of 0’s and 1’s as count_0 and count_1.

• We will traverse the array using for loop from i=0 to I <size.

• If the current element is 0 increment count_0.

• If not, calculate all possible subarrays with count_0 number of 0’s as temp_one_0=count*(count+1)/2.

• Add this to the previous temp_0 for subarrays with all 0’s found so far.

• Do a similar process using for loop for 1’s with variables as count_1, temp_one_1 and temp_1.

• At the end of both traversals, temp_0 and temp_1 will have respective counts of all subarrays within arr that have all 0’s and all 1’s.

Example

 Live Demo

#include <bits/stdc++.h>
using namespace std;
void sub_zero_one(int arr[], int size){
   int count_1 = 0;
   int count_0 = 0;
   int temp_1 = 0;
   int temp_0 = 0;
   for (int i = 0; i < size; i++){
      if (arr[i] == 1){
         count_1++;
      }
      else{
         int temp_one_1 = (count_1) * (count_1 + 1) / 2;
         temp_1 = temp_1 + temp_one_1;
         count_1 = 0;
      }
   }
   for (int i = 0; i < size; i++){
      if (arr[i] == 0)
         { count_0++; }
      else{
         int temp_one_0 = (count_0) * (count_0 + 1) / 2;
         temp_0 = temp_0 + temp_one_0;
         count_0 = 0;
      }
   }
   if (count_1){
      int temp_one_1 = (count_1) * (count_1 + 1) / 2;
      temp_1 = temp_1 + temp_one_1;
   }
   if (count_0){
      int temp_one_0 = (count_0) * (count_0 + 1) / 2;
      temp_0 = temp_0 + temp_one_0;
   }
   cout<<"Subarrays with only 0's are : "<<temp_0;
   cout<<"\nSubarrays with only 1's are : "<<temp_1;
}
int main(){
   int arr[] = { 0, 0, 0, 1, 1, 0, 1};
   int size = sizeof(arr) / sizeof(arr[0]);
   sub_zero_one(arr, size);
   return 0;
}

Output

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

Subarrays with only 0's are : 7
Subarrays with only 1's are : 4