Skip to content
Related Articles

Related Articles

Improve Article

Count subarrays having an equal count of 0s and 1s segregated

  • Difficulty Level : Easy
  • Last Updated : 05 May, 2021

Given a binary array arr[], the task is to count the number of subarrays having equal count of 0s and 1s, and all the 0s and 1s are placed consecutively in that subarray.

Examples:

Input: arr[] = {1, 0, 1, 1}
Output: 2
Explanation: The subarrays satisfying the given conditions are {1, 0} and {0, 1}. Therefore, the count of such subarrays is 2.

Input: arr[] = {1, 1, 0, 0, 1, 0}
Output: 4
Explanation: The subarrays satisfying the given conditions are {1, 1, 0, 0}, {1, 0}, {0, 1}, {1, 0}. Therefore, the count of such subarrays is 4.

Naive Approach: The simplest approach is to traverse the given array and for every pair of unequal adjacent elements, iterate the left and right of the current index and check if the count of 1s and 0s are equal or not. Increment the count of subarrays until found to be false. After complete traversal of the array, print the total count of subarrays.



Below is the implementation of the above approach:

C++




// C++ program for the above approach
 
#include <bits/stdc++.h>
using namespace std;
 
// Function to count subarrays
// having equal count of 0s and 1s
// with all 0s and all 1s grouped together
void countSubarrays(int A[], int N)
{
    // Stores the count of subarrays
    int ans = 0;
 
    for (int i = 0; i < N - 1; i++) {
 
        // If current element is different
        // from the next array element
        if (A[i] != A[i + 1]) {
 
            // Increment count
            ans++;
 
            // Count the frequency of
            // 1s and 0s
            for (int j = i - 1, k = i + 2;
                 j >= 0 && k < N
                 && A[j] == A[i]
                 && A[k] == A[i + 1];
                 j--, k++) {
 
                // Increment count
                ans++;
            }
        }
    }
 
    // Print the final count
    cout << ans << "\n";
}
 
// Driver Code
int main()
{
    int A[] = { 1, 1, 0, 0, 1, 0 };
    int N = sizeof(A) / sizeof(A[0]);
 
    // Function Call
    countSubarrays(A, N);
 
    return 0;
}

Java




// Java program for the above approach
import java.util.*;
 
class GFG
{
 
// Function to count subarrays
// having equal count of 0s and 1s
// with all 0s and all 1s grouped together
static void countSubarrays(int A[], int N)
{
   
    // Stores the count of subarrays
    int ans = 0;
 
    for (int i = 0; i < N - 1; i++)
    {
 
        // If current element is different
        // from the next array element
        if (A[i] != A[i + 1])
        {
 
            // Increment count
            ans++;
 
            // Count the frequency of
            // 1s and 0s
            for (int j = i - 1, k = i + 2;
                 j >= 0 && k < N
                 && A[j] == A[i]
                 && A[k] == A[i + 1];
                 j--, k++)
            {
 
                // Increment count
                ans++;
            }
        }
    }
 
    // Print the final count
    System.out.print(ans+ "\n");
}
 
// Driver Code
public static void main(String[] args)
{
    int A[] = { 1, 1, 0, 0, 1, 0 };
    int N = A.length;
 
    // Function Call
    countSubarrays(A, N);
}
}
 
// This code is contributed by shikhasingrajput

Python3




# Python3 program for the above approach
 
# Function to count subarrays
# having equal count of 0s and 1s
# with all 0s and all 1s grouped together
def countSubarrays(A, N) :
 
    # Stores the count of subarrays
    ans = 0;
 
    for i in range(N - 1) :
 
        # If current element is different
        # from the next array element
        if (A[i] != A[i + 1]) :
 
            # Increment count
            ans += 1;
 
            # Count the frequency of
            # 1s and 0s
            j = i - 1; k = i + 2;
            while (j >= 0 and k < N and A[j] == A[i] and A[k] == A[i + 1]) :
                 
                # Increment count
                ans += 1;
                 
                j -= 1;
                k += 1;
 
    # Print the final count
    print(ans);
 
# Driver Code
if __name__ == "__main__" :
 
    A = [ 1, 1, 0, 0, 1, 0 ];
    N = len(A);
 
    # Function Call
    countSubarrays(A, N);
 
    # This code is contributed by AnkitRai01

C#




// C# program for the above approach
using System;
  
class GFG{
  
// Function to count subarrays
// having equal count of 0s and 1s
// with all 0s and all 1s grouped together
static void countSubarrays(int[] A, int N)
{
     
    // Stores the count of subarrays
    int ans = 0;
  
    for(int i = 0; i < N - 1; i++)
    {
         
        // If current element is different
        // from the next array element
        if (A[i] != A[i + 1])
        {
  
            // Increment count
            ans++;
  
            // Count the frequency of
            // 1s and 0s
            for(int j = i - 1, k = i + 2;
                    j >= 0 && k < N &&
                       A[j] == A[i] &&
                       A[k] == A[i + 1];
                    j--, k++)
            {
                 
                // Increment count
                ans++;
            }
        }
    }
  
    // Print the final count
    Console.Write(ans + "\n");
}
  
// Driver Code
public static void Main()
{
    int[] A = { 1, 1, 0, 0, 1, 0 };
    int N = A.Length;
  
    // Function Call
    countSubarrays(A, N);
}
}
 
// This code is contributed by sanjoy_62

Javascript




<script>
// JavaScript program to implement
// the above approach
 
// Function to count subarrays
// having equal count of 0s and 1s
// with all 0s and all 1s grouped together
function countSubarrays(A, N)
{
    
    // Stores the count of subarrays
    let ans = 0;
  
    for (let i = 0; i < N - 1; i++)
    {
  
        // If current element is different
        // from the next array element
        if (A[i] != A[i + 1])
        {
  
            // Increment count
            ans++;
  
            // Count the frequency of
            // 1s and 0s
            for (let j = i - 1, k = i + 2;
                 j >= 0 && k < N
                 && A[j] == A[i]
                 && A[k] == A[i + 1];
                 j--, k++)
            {
  
                // Increment count
                ans++;
            }
        }
    }
  
    // Prlet the final count
    document.write(ans+ "<br/>");
}
 
 
// Driver Code
 
    let A = [ 1, 1, 0, 0, 1, 0 ];
    let N = A.length;
  
    // Function Call
    countSubarrays(A, N);
     
</script>

 
 

Output: 
4

 

 

Time Complexity: O(N2)
Auxiliary Space: O(1)

 

Efficient Approach: To optimize the above approach, follow the steps below:

 

  • Initialize a variable, say res, to store the count of subarrays.
  • Initialize a variable, say curr, with the first value of the array and an array cnt[], to keep track of the consecutive elements.
  • Traverse the array and perform the following steps:
    • If the current element is equal to curr, increment the last value of cnt[].
    • Otherwise, update curr to the current element and append 1 to the array cnt[].
  • Traverse the array cnt[] and find the sum of the minimum of the adjacent elements and add it to the variable res. This ensures that the frequency of the elements is equal.
  • After completing the above steps, print the value of res as the resultant count of subarrays.

 

Below is the implementation of the above approach:

 

C++




// C++ program for the above approach
#include <bits/stdc++.h>
using namespace std;
 
// Function to count subarrays
// having equal count of 0s and 1s
// with all 0s and all 1s grouped together
void countSubarrays(int A[], int N)
{
 
    // Stores the count
    int res = 0;
 
    // Initialize cur with first element
    int curr = A[0];
    vector<int> cnt = {1};
    for (int c = 1; c < N; c++)
    {
 
        // If the next element is same
        // as the current element
        if (A == curr)
 
            // Increment count
            cnt[cnt.size() - 1]++;
        else
 
            // Update curr
            curr = A;
        cnt.push_back(1);
    }
 
    // Iterate over the array count
    for (int i = 1; i < cnt.size(); i++)
    {
 
        // Consider the minimum
        res += min(cnt[i - 1], cnt[i]);
    }
    cout << (res - 1);
}
 
// Driver code
int main()
{
    // Given arr[]
    int A[] = { 1, 1, 0, 0, 1, 0 };   
    int N = sizeof(A) / sizeof(A[0]);
   
    // Function Call
    countSubarrays(A, N);
    return 0;
}
 
// This code is contributed by divyesh072019

Java




import java.util.Vector;
 
// Java program for the above approach
class GFG {
   
    // Function to count subarrays
    // having equal count of 0s and 1s
    // with all 0s and all 1s grouped together
    static void countSubarrays(int[] A)
    {
 
        // Stores the count
        int res = 0;
 
        // Initialize cur with first element
        int curr = A[0];
        int[] cnt = new int[A.length];
        cnt[0] = 1;
        for (int c = 1; c < A.length; c++) {
 
            // If the next element is same
            // as the current element
            if (A == curr)
 
                // Increment count
 
                cnt++;
            else
 
                // Update curr
                curr = A;
            cnt = 1;
        }
 
        // Iterate over the array count
        for (int i = 1; i < cnt.length; i++) {
 
            // Consider the minimum
            res += Math.min(cnt[i - 1], cnt[i]);
        }
        System.out.println(res - 1);
    }
   
    // Driver code
    public static void main(String[] args)
    {
        // Given arr[]
        int[] A = { 1, 1, 0, 0, 1, 0 };
 
        // Function Call
        countSubarrays(A);
    }
}
 
// This code is contributed by 29AjayKumar

Python3




# Python3 program for the above approach
 
# Function to count subarrays
# having equal count of 0s and 1s
# with all 0s and all 1s grouped together
def countSubarrays(A):
     
    # Stores the count
    res = 0
     
    # Initialize cur with first element
    curr, cnt = A[0], [1]
    for c in A[1:]:
       
        # If the next element is same
        # as the current element
        if c == curr:
           
            # Increment count
            cnt[-1] += 1
        else:
           
            # Update curr
            curr = c
        cnt.append(1)
     
    # Iterate over the array count
    for i in range(1, len(cnt)):
       
        # Consider the minimum
        res += min(cnt[i - 1], cnt[i])
     
    print(res - 1)
 
# Given arr[]
A = [1, 1, 0, 0, 1, 0]
 
# Function Call
countSubarrays(A)

C#




// C# program for the above approach
using System;
 
class GFG{
   
// Function to count subarrays
// having equal count of 0s and 1s
// with all 0s and all 1s grouped together
static void countSubarrays(int[] A)
{
     
    // Stores the count
    int res = 0;
 
    // Initialize cur with first element
    int curr = A[0];
    int[] cnt = new int[A.Length];
    cnt[0] = 1;
     
    for(int c = 1; c < A.Length; c++)
    {
         
        // If the next element is same
        // as the current element
        if (A == curr)
 
            // Increment count
            cnt++;
        else
 
            // Update curr
            curr = A;
             
        cnt = 1;
    }
 
    // Iterate over the array count
    for(int i = 1; i < cnt.Length; i++)
    {
         
        // Consider the minimum
        res += Math.Min(cnt[i - 1], cnt[i]);
    }
    Console.WriteLine(res - 1);
}
 
// Driver code
public static void Main(String[] args)
{
     
    // Given []arr
    int[] A = { 1, 1, 0, 0, 1, 0 };
 
    // Function Call
    countSubarrays(A);
}
}
 
// This code is contributed by Amit Katiyar

Javascript




<script>
// Javascript program for the above approach
 
// Function to count subarrays
// having equal count of 0s and 1s
// with all 0s and all 1s grouped together
function countSubarrays( A, N)
{
 
    // Stores the count
    var res = 0;
 
    // Initialize cur with first element
    var curr = A[0];
    var cnt = [];
    cnt.fill(1)
    for (var c = 1; c < N; c++)
    {
 
        // If the next element is same
        // as the current element
        if (A == curr)
 
            // Increment count
            cnt[cnt.length - 1]++;
        else
 
            // Update curr
            curr = A;
        cnt.push(1);
    }
 
    // Iterate over the array count
    for (var i = 1; i < cnt.length; i++)
    {
 
        // Consider the minimum
        res += Math.min(cnt[i - 1], cnt[i]);
    }
    document.write (res);
}
 
var A = [ 1, 1, 0, 0, 1, 0 ];   
var N = A.length;
   
    // Function Call
    countSubarrays(A, N);
 
// This code is contributed by SoumikMondal
</script>
Output: 
4

 

Time Complexity: O(N)
Auxiliary Space: O(N)

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.  To complete your preparation from learning a language to DS Algo and many more,  please refer Complete Interview Preparation Course.

In case you wish to attend live classes with experts, please refer DSA Live Classes for Working Professionals and Competitive Programming Live for Students.




My Personal Notes arrow_drop_up
Recommended Articles
Page :