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

24 Mar, 2022
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
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)
The problem is closely related to the Largest subarray with an equal number of 0’s and 1’sApproach: Following are the steps: 
Consider all 0’s in arr[] as -1.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.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.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.Also, add freq of the sum of 0 to the hash table for the final result.
Consider all 0’s in arr[] as -1.
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.
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.
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.
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. 
C++
Java
Python3
C#
Javascript
// 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'sint 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 aboveint 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 implementation to count subarrays with// equal number of 1's and 0'simport java.util.*; class GFG{ // function to count subarrays with// equal number of 1's and 0'sstatic 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 abovepublic 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 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'sdef 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 abovearr = [ 1, 0, 0, 1, 0, 1, 1 ]n = len(arr)print("Count =",    countSubarrWithEqualZeroAndOne(arr, n)) # This code is contributed by "Sharad_Bhardwaj".
// C# implementation to count subarrays// with equal number of 1's and 0'susing System;using System.Collections.Generic; class GFG{ // function to count subarrays with// equal number of 1's and 0'sstatic 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 abovepublic 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
<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'sfunction 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 abovevar 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>
Output:  
Count = 8
Time Complexity: O(n). Auxiliary Space: O(n).  
Another approach:  This approach is very easy to understand if we simply know this question : Number of subarrays having sum exactly equal to k.  We just have to make k=0 and all the zeroes in the given array to -1. 
C++
Java
Python3
C#
Javascript
#include <bits/stdc++.h> using namespace std; int countSubarrWithEqualZeroAndOne(int arr[], int n){    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);}
import java.util.HashMap;import java.util.Map; // Java implementation to count subarrays with// equal number of 1's and 0'spublic 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 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 Codearr = [1, 0, 0, 1, 0, 1, 1] n = len(arr) print("count =",      countSubarrWithEqualZeroAndOne(arr, n)) # This code is contributed by mohit kumar
// C# implementation to count subarrays with// equal number of 1's and 0'susing 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
<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 abovevar 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>
Output:  
Count = 8
 
ravi_Sgo
shivamgoel009
mohit kumar 29
29AjayKumar
Rajput-Ji
princiraj1992
noob2000
shivanisinghss2110
dishantarora
superdude
binary-string
Arrays
Hash
Arrays
Hash
Writing code in comment?
Please use ide.geeksforgeeks.org,
generate link and share the link here.
Comments
Old Comments
Stack Data Structure (Introduction and Program)
Top 50 Array Coding Problems for Interviews
Introduction to Arrays
Multidimensional Arrays in Java
Linear Search
Given an array A[] and a number x, check for pair in A[] with sum as x (aka Two Sum)
Internal Working of HashMap in Java
Hashing | Set 1 (Introduction)
Count pairs with given sum
Hashing | Set 3 (Open Addressing)