Open In App

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

Last Updated : 21 Nov, 2021
Improve
Improve
Like Article
Like
Save
Share
Report

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++;
            }
        }
    }
  
    // Print 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)



Similar Reads

Number of ways to select equal sized subarrays from two arrays having atleast K equal pairs of elements
Given two arrays A[] and B[], and an integer K, the task is to find the number of ways to select two subarrays of the same size, one from A and the other one from B such that the subarrays have at least K equal pairs of elements. (i.e. the number of pairs (A[i], B[j]) in the two selected subarrays such that A[i] = B[j] &gt;= K). Examples: Input: A[
13 min read
Split an Array to maximize subarrays having equal count of odd and even elements for a cost not exceeding K
Given an array arr[] of size N and an integer K, the task is to split the given array into maximum possible subarrays having equal count of even and odd elements such that the cost to split the array does not exceed K. The cost to split an array into a subarray is the difference between the last and first elements of the subarrays respectively. Exa
8 min read
Count subarrays having sum of elements at even and odd positions equal
Given an array arr[] of integers, the task is to find the total count of subarrays such that the sum of elements at even position and sum of elements at the odd positions are equal. Examples: Input: arr[] = {1, 2, 3, 4, 1}Output: 1Explanation: {3, 4, 1} is the only subarray in which sum of elements at even position {3, 1} = sum of element at odd po
7 min read
Count distinct sequences obtained by replacing all elements of subarrays having equal first and last elements with the first element any number of times
Given an array arr[] consisting of N integers, the task is to find the number of different sequences that can be formed after performing the below operation on the given array arr[] any number of times. Choose two indices i and j such that arr[i] is equal to arr[j] and update all the elements in the range [i, j] in the array to arr[i]. Examples: In
7 min read
Count of subarrays having sum equal to its length
Given an array arr[] of size N, the task is to find the number of subarrays having the sum of its elements equal to the number of elements in it. Examples: Input: N = 3, arr[] = {1, 0, 2}Output: 3Explanation:Total number of subarrays are 6 i.e., {1}, {0}, {2}, {1, 0}, {0, 2}, {1, 0, 2}.Out of 6 only three subarrays have the number of elements equal
7 min read
Count even length subarrays having bitwise XOR equal to 0
Given an array arr[] of size N, the task is to count all possible even length subarrays having bitwise XOR of subarray elements equal to 0. Examples: Input: arr[] = {2, 2, 3, 3, 6, 7, 8}Output: 3Explanation:Subarrays having XOR of elements equal to 0 are: {{2, 2}, {3, 3}, {2, 2, 3, 3}}Therefore, the required output is 3. Input: arr[] = {1, 2, 3, 3}
13 min read
Count of subarrays having sum equal to its length | Set 2
Given an array arr[] of size N, the task is to find the number of subarrays having sum of its elements equal to the number of elements in it. Examples: Input: N = 3, arr[] = {1, 0, 2}Output: 3Explanation:Total number of subarrays are 6 i.e., {1}, {0}, {2}, {1, 0}, {0, 2}, {1, 0, 2}.Out of the 6 subarrays, following three subarrays satisfy the given
7 min read
Count subarrays having product equal to the power of a given Prime Number
Given an array arr[] of size N and an integer M, the task is to count the number of subarrays having product of its elements equal to the power of M, where M is a prime number. Examples: Input: arr[] = {2, 2, 2, 2}, M = 2Output: 10Explanation: All possible non-empty subarrays having product equal to the power of M = (4 * (4 + 1)) / 2 = 10 Input: ar
8 min read
Count ways to split array into three non-empty subarrays having equal Bitwise XOR values
Given an array arr[] consisting of N non-negative integers, the task is to count the number of ways to split the array into three different non-empty subarrays such that Bitwise XOR of each subarray is equal. Examples: Input: arr[] = {7, 0, 5, 2, 7} Output: 2Explanation: All possible ways are:{{7}, {0, 5, 2}, {7}} where XOR value of each subarray i
9 min read
Count of subarrays in range [L, R] having XOR + 1 equal to XOR (XOR) 1 for M queries
Given an array, arr[] of N positive integers and M queries which consist of two integers [Li, Ri] where 1 ? Li ? Ri ? N. For each query, find the number of subarrays in range [Li, Ri] for which (X+1)=(X?1) where X denotes the xor of a subarray. Input: arr[]= {1, 2, 9, 8, 7}, queries[] = {{1, 5}, {3, 4}}Output: 6 1Explanation: Query 1: L=1, R=5: sub
12 min read