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)
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; } |
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". |
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
Output:
Count = 8
leave a comment
0 Comments