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

Given an array arr[] of size n containing 0 and 1 only. The problem is to count the subarrays having 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)

Recommended: Please solve it on PRACTICE first, before moving on to the solution.

The problem is closely related to Largest subarray with equal number of 0’s and 1’s.

Approach: 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 cumulative sum and then we get increment count by 1 for that sum represented as index in the hash table. Sub-array by each pair of positions with same value of cumulative sum constitute a continuous range with 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 pair of indices of sub-array by (freq * (freq – 1)) / 2 number of ways . Do the same for all freq and sum up the result that will be the number all possible sub-arrays containing equal number of 1’s and 0’s.
Also add freq of the sum 0 in the hash table to 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 equal number of 1’s and 0’s.

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  

Output:

Count = 8

Time Complexity: O(n).
Auxiliary Space: O(n).
Another approach:

C++

#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]) 
                count += mp[sum]; 
            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 

Output:

Count = 8

My Personal Notes

Oct	NOV	Dec
	22
2018	2019	2020

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

Recommended: Please solve it on PRACTICE first, before moving on to the solution.

C++

Java

Python3

C#

C++

Java

Python3

C#

Recommended Posts:

Suggest a Topic

Recommended: Please solve it on PRACTICE first, before moving on to the solution.

C++

Java

Python3

C#

C++

Java

Python3

C#

Recommended Posts: