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

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# 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:

Time Complexity: O(N²)
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# 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:

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

Article Tags :

Arrays

DSA

Greedy

Searching

frequency-counting

subarray