Open In App

Count subarrays with equal number of 1's and 0's

Last Updated : 29 Nov, 2023
Comments
Improve
Suggest changes
Like Article
Like
Report
Try it on GfG Practice
redirect icon

Given an array arr[] of size n containing 0 and 1 only. The problem is to count the subarrays having an equal number of 0's and 1's.

Examples:  

Input: arr[] = {1, 0, 0, 1, 0, 1, 1}
Output: 8
Explanation: The index range for the 8 sub-arrays are: (0, 1), (2, 3), (0, 3), (3, 4), (4, 5)(2, 5), (0, 5), (1, 6)

Input: arr = { 1, 0, 0, 1, 1, 0, 0, 1}
Output: 12

Count subarrays with equal number of 1's and 0's using Frequency Counting:

The problem is closely related to the Largest subarray with an equal number of 0's and 1's

Follow the steps below to implement the above idea:

  1. Consider all 0's in arr[] as -1.
  2. Create a hash table that holds the count of each sum[i] value, where sum[i] = sum(arr[0]+..+arr[i]), for i = 0 to n-1.
  3. Now start calculating the cumulative sum and then we get an incremental count of 1 for that sum represented as an index in the hash table. Arrays of each pair of positions with the same value in the cumulative sum constitute a continuous range with an equal number of 1's and 0's.
  4. Now traverse the hash table and get the frequency of each element in the hash table. Let frequency be denoted as freq. For each freq > 1 we can choose any two pairs of indices of a sub-array by (freq * (freq – 1)) / 2 number of ways. Do the same for all freq and sum up the result will be the number of all possible sub-arrays containing the equal number of 1's and 0's.
  5. Also, add freq of the sum of 0 to the hash table for the final result.

Explanation: 
Considering all 0’s as -1. if sum[i] == sum[j], where sum[i] = sum(arr[0]+..+arr[i]) and sum[j] = sum(arr[0]+..+arr[j]) and ‘i’ is less than ‘j’, then sum(arr[i+1]+..+arr[j]) must be 0. It can only be 0 if arr(i+1, .., j) contains an equal number of 1’s and 0’s. 

Follow the steps below to implement the above approach:

C++ 
// C++ implementation to count subarrays with
// equal number of 1's and 0's
#include <bits/stdc++.h>

using namespace std;

// function to count subarrays with
// equal number of 1's and 0's
int countSubarrWithEqualZeroAndOne(int arr[], int n)
{
    // 'um' implemented as hash table to store
    // frequency of values obtained through
    // cumulative sum
    unordered_map<int, int> um;
    int curr_sum = 0;

    // Traverse original array and compute cumulative
    // sum and increase count by 1 for this sum
    // in 'um'. Adds '-1' when arr[i] == 0
    for (int i = 0; i < n; i++) {
        curr_sum += (arr[i] == 0) ? -1 : arr[i];
        um[curr_sum]++;
    }

    int count = 0;
    // traverse the hash table 'um'
    for (auto itr = um.begin(); itr != um.end(); itr++) {

        // If there are more than one prefix subarrays
        // with a particular sum
        if (itr->second > 1)
            count
                += ((itr->second * (itr->second - 1)) / 2);
    }

    // add the subarrays starting from 1st element and
    // have equal number of 1's and 0's
    if (um.find(0) != um.end())
        count += um[0];

    // required count of subarrays
    return count;
}

// Driver program to test above
int main()
{
    int arr[] = { 1, 0, 0, 1, 0, 1, 1 };
    int n = sizeof(arr) / sizeof(arr[0]);
    cout << "Count = "
         << countSubarrWithEqualZeroAndOne(arr, n);
    return 0;
}
Java 
// Java implementation to count subarrays with
// equal number of 1's and 0's
import java.util.*;

class GFG {

    // function to count subarrays with
    // equal number of 1's and 0's
    static int countSubarrWithEqualZeroAndOne(int arr[],
                                              int n)
    {
        // 'um' implemented as hash table to store
        // frequency of values obtained through
        // cumulative sum
        Map<Integer, Integer> um = new HashMap<>();
        int curr_sum = 0;

        // Traverse original array and compute cumulative
        // sum and increase count by 1 for this sum
        // in 'um'. Adds '-1' when arr[i] == 0
        for (int i = 0; i < n; i++) {
            curr_sum += (arr[i] == 0) ? -1 : arr[i];
            um.put(curr_sum, um.get(curr_sum) == null
                                 ? 1
                                 : um.get(curr_sum) + 1);
        }

        int count = 0;

        // traverse the hash table 'um'
        for (Map.Entry<Integer, Integer> itr :
             um.entrySet()) {

            // If there are more than one prefix subarrays
            // with a particular sum
            if (itr.getValue() > 1)
                count += ((itr.getValue()
                           * (itr.getValue() - 1))
                          / 2);
        }

        // add the subarrays starting from 1st element and
        // have equal number of 1's and 0's
        if (um.containsKey(0))
            count += um.get(0);

        // required count of subarrays
        return count;
    }

    // Driver program to test above
    public static void main(String[] args)
    {
        int arr[] = { 1, 0, 0, 1, 0, 1, 1 };
        int n = arr.length;
        System.out.println(
            "Count = "
            + countSubarrWithEqualZeroAndOne(arr, n));
    }
}

// This code is contributed by Rajput-Ji
Python3 
# Python3 implementation to count
# subarrays with equal number
# of 1's and 0's

# function to count subarrays with
# equal number of 1's and 0's


def countSubarrWithEqualZeroAndOne(arr, n):

    # 'um' implemented as hash table
    # to store frequency of values
    # obtained through cumulative sum
    um = dict()
    curr_sum = 0

    # Traverse original array and compute
    # cumulative sum and increase count
    # by 1 for this sum in 'um'.
    # Adds '-1' when arr[i] == 0
    for i in range(n):
        curr_sum += (-1 if (arr[i] == 0) else arr[i])
        if um.get(curr_sum):
            um[curr_sum] += 1
        else:
            um[curr_sum] = 1

    count = 0

    # traverse the hash table 'um'
    for itr in um:

        # If there are more than one
        # prefix subarrays with a
        # particular sum
        if um[itr] > 1:
            count += ((um[itr] * int(um[itr] - 1)) / 2)

    # add the subarrays starting from
    # 1st element and have equal
    # number of 1's and 0's
    if um.get(0):
        count += um[0]

    # required count of subarrays
    return int(count)


# Driver code to test above
arr = [1, 0, 0, 1, 0, 1, 1]
n = len(arr)
print("Count =",
      countSubarrWithEqualZeroAndOne(arr, n))

# This code is contributed by "Sharad_Bhardwaj".
C# 
// C# implementation to count subarrays
// with equal number of 1's and 0's
using System;
using System.Collections.Generic;

class GFG {

    // function to count subarrays with
    // equal number of 1's and 0's
    static int countSubarrWithEqualZeroAndOne(int[] arr,
                                              int n)
    {
        // 'um' implemented as hash table to store
        // frequency of values obtained through
        // cumulative sum
        Dictionary<int, int> mp
            = new Dictionary<int, int>();
        int curr_sum = 0;

        // Traverse original array and compute cumulative
        // sum and increase count by 1 for this sum
        // in 'um'. Adds '-1' when arr[i] == 0
        for (int i = 0; i < n; i++) {
            curr_sum += (arr[i] == 0) ? -1 : arr[i];
            if (mp.ContainsKey(curr_sum)) {
                var v = mp[curr_sum];
                mp.Remove(curr_sum);
                mp.Add(curr_sum, ++v);
            }
            else
                mp.Add(curr_sum, 1);
        }

        int count = 0;

        // traverse the hash table 'um'
        foreach(KeyValuePair<int, int> itr in mp)
        {

            // If there are more than one prefix subarrays
            // with a particular sum
            if (itr.Value > 1)
                count
                    += ((itr.Value * (itr.Value - 1)) / 2);
        }

        // add the subarrays starting from 1st element
        // and have equal number of 1's and 0's
        if (mp.ContainsKey(0))
            count += mp[0];

        // required count of subarrays
        return count;
    }

    // Driver program to test above
    public static void Main(String[] args)
    {
        int[] arr = { 1, 0, 0, 1, 0, 1, 1 };
        int n = arr.Length;
        Console.WriteLine(
            "Count = "
            + countSubarrWithEqualZeroAndOne(arr, n));
    }
}

// This code is contributed by PrinciRaj1992
JavaScript 
<script>

// Javascript implementation to count subarrays with
// equal number of 1's and 0's

// function to count subarrays with
// equal number of 1's and 0's
function countSubarrWithEqualZeroAndOne(arr, n)
{
    // 'um' implemented as hash table to store
    // frequency of values obtained through
    // cumulative sum
    var um = new Map();
    var curr_sum = 0;

    // Traverse original array and compute cumulative
    // sum and increase count by 1 for this sum
    // in 'um'. Adds '-1' when arr[i] == 0
    for (var i = 0; i < n; i++) {
        curr_sum += (arr[i] == 0) ? -1 : arr[i];
        if(um.has(curr_sum))
            um.set(curr_sum, um.get(curr_sum)+1);
        else
            um.set(curr_sum, 1)
    }

    var count = 0;
    // traverse the hash table 'um'
    um.forEach((value, key) => {
        
        // If there are more than one prefix subarrays
        // with a particular sum
        if (value > 1)
            count += ((value * (value - 1)) / 2);
    });

    // add the subarrays starting from 1st element and
    // have equal number of 1's and 0's
    if (um.has(0))
        count += um.get(0);

    // required count of subarrays
    return count;
}

// Driver program to test above
var arr = [1, 0, 0, 1, 0, 1, 1];
var n = arr.length;
document.write( "Count = "
      + countSubarrWithEqualZeroAndOne(arr, n));

// This code is contributed by noob2000.
</script>
PHP 
<?php
// Php program for the above approach
// Function to count subarrays with equal number of 1's and 0's
function countSubarrWithEqualZeroAndOne($arr, $n) {
    // 'um' implemented as hash table to store
    // frequency of values obtained through
    // cumulative sum
    $um = array();
    $curr_sum = 0;

    // Traverse original array and compute cumulative
    // sum and increase count by 1 for this sum
    // in 'um'. Adds '-1' when $arr[$i] == 0
    for ($i = 0; $i < $n; $i++) {
        $curr_sum += ($arr[$i] == 0) ? -1 : $arr[$i];
        $um[$curr_sum]++;
    }

    $count = 0;

    // Traverse the hash table 'um'
    foreach ($um as $value) {
        // If there are more than one prefix subarrays
        // with a particular sum
        if ($value > 1) {
            $count += (($value * ($value - 1)) / 2);
        }
    }

    // Add the subarrays starting from 1st element and
    // have an equal number of 1's and 0's
    if (array_key_exists(0, $um)) {
        $count += $um[0];
    }

    // Required count of subarrays
    return $count;
}

// Driver program to test above
$arr = array(1, 0, 0, 1, 0, 1, 1);
$n = count($arr);
echo "Count = " . countSubarrWithEqualZeroAndOne($arr, $n);

?>

Output
Count = 8

Time Complexity: O(N), where N is the length of the given array
Auxiliary Space: O(N). 

Count subarrays with equal number of 1's and 0's using Map:

Follow the steps below for implementation:

  • Create a map say mp.
  • Iterate over the length of the given array
    • Check if arr[i] == 0, then replace it with -1.
    • Keep calculating the number into sum till ith index.
    • If sum = 0, it implies the number of 0's and 1's are equal from arr[0]..arr[i]
    • if mp[sum] exists then add "frequency-1" to count
    • if the frequency of "sum" is zero then we initialize that frequency to 1, f it's not 0, we increment it
  • Finally, return the count.

Follow the steps below to implement the above approach:

C++ 
#include <bits/stdc++.h>

using namespace std;

int countSubarrWithEqualZeroAndOne(int arr[], int n)
{
    unordered_map<int, int> mp;
    int sum = 0;
    int count = 0;
    for (int i = 0; i < n; i++) {
        // Replacing 0's in array with -1
        if (arr[i] == 0)
            arr[i] = -1;

        sum += arr[i];

        // If sum = 0, it implies number of 0's and 1's are
        // equal from arr[0]..arr[i]
        if (sum == 0)
            count++;

        // if mp[sum] exists then add "frequency-1" to count
        if (mp[sum])
            count += mp[sum];

        // if frequency of "sum" is zero then we initialize
        // that frequency to 1 if its not 0, we increment it
        if (mp[sum] == 0)
            mp[sum] = 1;
        else
            mp[sum]++;
    }
    return count;
}

int main()
{
    int arr[] = { 1, 0, 0, 1, 0, 1, 1 };

    int n = sizeof(arr) / sizeof(arr[0]);

    cout << "count="
         << countSubarrWithEqualZeroAndOne(arr, n);
}
Java 
import java.util.HashMap;
import java.util.Map;

// Java implementation to count subarrays with
// equal number of 1's and 0's
public class Main {

    // Function that returns count of sub arrays
    // with equal numbers of 1's and 0's
    static int countSubarrWithEqualZeroAndOne(int[] arr,
                                              int n)
    {
        Map<Integer, Integer> myMap = new HashMap<>();
        int sum = 0;
        int count = 0;
        for (int i = 0; i < n; i++) {
            // Replacing 0's in array with -1
            if (arr[i] == 0)
                arr[i] = -1;

            sum += arr[i];

            // If sum = 0, it implies number of 0's and 1's
            // are equal from arr[0]..arr[i]
            if (sum == 0)
                count++;

            if (myMap.containsKey(sum))
                count += myMap.get(sum);

            if (!myMap.containsKey(sum))
                myMap.put(sum, 1);
            else
                myMap.put(sum, myMap.get(sum) + 1);
        }
        return count;
    }

    // main function
    public static void main(String[] args)
    {
        int arr[] = { 1, 0, 0, 1, 0, 1, 1 };
        int n = arr.length;
        System.out.println(
            "Count = "
            + countSubarrWithEqualZeroAndOne(arr, n));
    }
}
Python3 
# Python3 implementation to count subarrays
# with equal number of 1's and 0's


def countSubarrWithEqualZeroAndOne(arr, n):
    mp = dict()
    Sum = 0
    count = 0

    for i in range(n):

        # Replacing 0's in array with -1
        if (arr[i] == 0):
            arr[i] = -1

        Sum += arr[i]

        # If Sum = 0, it implies number of
        # 0's and 1's are equal from arr[0]..arr[i]
        if (Sum == 0):
            count += 1

        if (Sum in mp.keys()):
            count += mp[Sum]

        mp[Sum] = mp.get(Sum, 0) + 1

    return count


# Driver Code
arr = [1, 0, 0, 1, 0, 1, 1]

n = len(arr)

print("count =",
      countSubarrWithEqualZeroAndOne(arr, n))

# This code is contributed by mohit kumar
C# 
// C# implementation to count subarrays with
// equal number of 1's and 0's
using System;
using System.Collections.Generic;

class GFG {

    // Function that returns count of sub arrays
    // with equal numbers of 1's and 0's
    static int countSubarrWithEqualZeroAndOne(int[] arr,
                                              int n)
    {
        Dictionary<int, int> myMap
            = new Dictionary<int, int>();
        int sum = 0;
        int count = 0;
        for (int i = 0; i < n; i++) {
            // Replacing 0's in array with -1
            if (arr[i] == 0)
                arr[i] = -1;

            sum += arr[i];

            // If sum = 0, it implies number of 0's and 1's
            // are equal from arr[0]..arr[i]
            if (sum == 0)
                count++;

            if (myMap.ContainsKey(sum))
                count += myMap[sum];

            if (!myMap.ContainsKey(sum))
                myMap.Add(sum, 1);
            else {
                var v = myMap[sum] + 1;
                myMap.Remove(sum);
                myMap.Add(sum, v);
            }
        }
        return count;
    }

    // Driver code
    public static void Main(String[] args)
    {
        int[] arr = { 1, 0, 0, 1, 0, 1, 1 };
        int n = arr.Length;
        Console.WriteLine(
            "Count = "
            + countSubarrWithEqualZeroAndOne(arr, n));
    }
}

// This code is contributed by 29AjayKumar
JavaScript 
<script>
 
// Javascript implementation to count subarrays with
// equal number of 1's and 0's

function countSubarrWithEqualZeroAndOne(arr, n)
{
    var mp = new Map();
    var sum = 0;
    let count = 0;
    
    for (var i = 0; i < n; i++) {
        //Replacing 0's in array with -1
        if (arr[i] == 0)
            arr[i] = -1;
 
        sum += arr[i];
 
        //If sum = 0, it implies number of 0's and 1's are
        //equal from arr[0]..arr[i]
        if (sum == 0)
            count += 1;
 
        if (mp.has(sum) == true)
            count += mp.get(sum);
                
        if(mp.has(sum) == false)
            mp.set(sum, 1);
        else
            mp.set(sum, mp.get(sum)+1);
    }
      return count;
}
 
// Driver program to test above
var arr = [1, 0, 0, 1, 0, 1, 1];
var n = arr.length;
document.write( "Count = "
      + countSubarrWithEqualZeroAndOne(arr, n));
 
// This code is contributed by noob2000.
</script>
PHP 
<?php
//Php code for the above approach
// Function to count subarrays with equal number of 1's and 0's
function countSubarrWithEqualZeroAndOne($arr, $n) {
    // Initialize an associative array to store the cumulative sum frequencies
    $mp = array();
    $sum = 0;
    $count = 0;

    // Traverse the array
    for ($i = 0; $i < $n; $i++) {
        // Replacing 0's in the array with -1
        if ($arr[$i] == 0) {
            $arr[$i] = -1;
        }

        $sum += $arr[$i];

        // If sum = 0, it implies the number of 0's and 1's are
        // equal from $arr[0]..$arr[$i]
        if ($sum == 0) {
            $count++;
        }

        // If $mp[$sum] exists, then add "frequency-1" to count
        if (isset($mp[$sum])) {
            $count += $mp[$sum];
        }

        // If frequency of $sum is zero, then initialize that frequency to 1,
        // if it's not zero, increment it
        if (!isset($mp[$sum])) {
            $mp[$sum] = 1;
        } else {
            $mp[$sum]++;
        }
    }

    // Return the final count of subarrays
    return $count;
}

// Driver program to test above
$arr = array(1, 0, 0, 1, 0, 1, 1);
$n = count($arr);
echo "Count = " . countSubarrWithEqualZeroAndOne($arr, $n);

?>

Output
count=8

Time Complexity: O(N), where N is the length of the given array
Auxiliary Space: O(N).


Next Article
Count subarrays with equal number of 1's and 0's

A
ayushjauhari14
Improve
Article Tags :
Practice Tags :

Similar Reads

We use cookies to ensure you have the best browsing experience on our website. By using our site, you acknowledge that you have read and understood our Cookie Policy & Privacy Policy
Lightbox