Related Articles

# 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 ``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);` `    ``// 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

 ``

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 ``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;``    ``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);``  ` `    ``// 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;``    ``int``[] cnt = ``new` `int``[A.Length];``    ``cnt = 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

 ``
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