# Count Substrings with equal number of 0s, 1s and 2s

Given a string which consists of only 0s, 1s or 2s, count the number of substrings that have equal number of 0s, 1s and 2s.

**Examples:**

Input : str = “0102010” Output : 2 Explanation : Substring str[2, 4] = “102” and substring str[4, 6] = “201” has equal number of 0, 1 and 2 Input : str = "102100211" Output : 5

A **simple solution** is to iterate through all substring of str and checking whether they contain equal 0,1 and 2 or not. Total number of substring of str is O(n^{2}) checking each substring for count takes O(n) times, So total time taken to solve this problem is O(n^{3}) time with brute-force approach.

An **efficient solution** is to keep track of counts of 0, 1 and 2.

Let zc[i] denotes number of zeros between index 1 and i oc[i] denotes number of ones between index 1 and i tc[i] denotes number of twos between index 1 and i for substring str[i, j] to be counted in result we should have : zc[i] – zc[j-1] = oc[i] – oc[j-1] = tc[i] - tc[j-1] we can write above relation as follows : z[i] – o[i] = z[j-1] – o[j-1] and z[i] – t[i] = z[j-1] – t[j-1]

The above relations can be tracked while looping in string, at each index i we’ll calculate this difference pair and we’ll check how many time this difference pair has previously occurred and we’ll add that count to our result, for keeping track of previous difference pair we have used map in below code. Total time complexity of this solution is O(n log n) considering the fact that map operations, like search and insert take O(Log n) time.

## C++

`// C++ program to find substring with equal ` `// number of 0's, 1's and 2's ` `#include <bits/stdc++.h> ` `using` `namespace` `std; ` ` ` `// Method to count number of substring which ` `// has equal 0, 1 and 2 ` `int` `getSubstringWithEqual012(string str) ` `{ ` ` ` `int` `n = str.length(); ` ` ` ` ` `// map to store, how many times a difference ` ` ` `// pair has occurred previously ` ` ` `map< pair<` `int` `, ` `int` `>, ` `int` `> mp; ` ` ` `mp[make_pair(0, 0)] = 1; ` ` ` ` ` `// zc (Count of zeroes), oc(Count of 1s) ` ` ` `// and tc(count of twos) ` ` ` `// In starting all counts are zero ` ` ` `int` `zc = 0, oc = 0, tc = 0; ` ` ` ` ` `// looping into string ` ` ` `int` `res = 0; ` `// Initialize result ` ` ` `for` `(` `int` `i = 0; i < n; ++i) ` ` ` `{ ` ` ` `// increasing the count of current character ` ` ` `if` `(str[i] == ` `'0'` `) zc++; ` ` ` `else` `if` `(str[i] == ` `'1'` `) oc++; ` ` ` `else` `tc++; ` `// Assuming that string doesn't contain ` ` ` `// other characters ` ` ` ` ` `// making pair of differences (z[i] - o[i], ` ` ` `// z[i] - t[i]) ` ` ` `pair<` `int` `, ` `int` `> tmp = make_pair(zc - oc, ` ` ` `zc - tc); ` ` ` ` ` `// Count of previous occurrences of above pair ` ` ` `// indicates that the subarrays forming from ` ` ` `// every previous occurrence to this occurrence ` ` ` `// is a subarray with equal number of 0's, 1's ` ` ` `// and 2's ` ` ` `res = res + mp[tmp]; ` ` ` ` ` `// increasing the count of current difference ` ` ` `// pair by 1 ` ` ` `mp[tmp]++; ` ` ` `} ` ` ` ` ` `return` `res; ` `} ` ` ` `// driver code to test above method ` `int` `main() ` `{ ` ` ` `string str = ` `"0102010"` `; ` ` ` `cout << getSubstringWithEqual012(str) << endl; ` ` ` `return` `0; ` `} ` |

*chevron_right*

*filter_none*

## Python3

`# Python3 program to find substring with equal ` `# number of 0's, 1's and 2's ` ` ` `# Method to count number of substring which ` `# has equal 0, 1 and 2 ` `def` `getSubstringWithEqual012(string): ` ` ` `n ` `=` `len` `(string) ` ` ` ` ` `# map to store, how many times a difference ` ` ` `# pair has occurred previously ` ` ` `mp ` `=` `dict` `() ` ` ` `mp[(` `0` `, ` `0` `)] ` `=` `1` ` ` ` ` `# zc (Count of zeroes), oc(Count of 1s) ` ` ` `# and tc(count of twos) ` ` ` `# In starting all counts are zero ` ` ` `zc, oc, tc ` `=` `0` `, ` `0` `, ` `0` ` ` ` ` `# looping into string ` ` ` `res ` `=` `0` `# Initialize result ` ` ` `for` `i ` `in` `range` `(n): ` ` ` ` ` `# increasing the count of current character ` ` ` `if` `string[i] ` `=` `=` `'0'` `: ` ` ` `zc ` `+` `=` `1` ` ` `elif` `string[i] ` `=` `=` `'1'` `: ` ` ` `oc ` `+` `=` `1` ` ` `else` `: ` ` ` `tc ` `+` `=` `1` `# Assuming that string doesn't contain ` ` ` `# other characters ` ` ` ` ` `# making pair of differences (z[i] - o[i], ` ` ` `# z[i] - t[i]) ` ` ` `tmp ` `=` `(zc ` `-` `oc, zc ` `-` `tc) ` ` ` ` ` `# Count of previous occurrences of above pair ` ` ` `# indicates that the subarrays forming from ` ` ` `# every previous occurrence to this occurrence ` ` ` `# is a subarray with equal number of 0's, 1's ` ` ` `# and 2's ` ` ` `if` `tmp ` `not` `in` `mp: ` ` ` `res ` `+` `=` `0` ` ` `else` `: ` ` ` `res ` `+` `=` `mp[tmp] ` ` ` ` ` `# increasing the count of current difference ` ` ` `# pair by 1 ` ` ` `if` `tmp ` `in` `mp: ` ` ` `mp[tmp] ` `+` `=` `1` ` ` `else` `: ` ` ` `mp[tmp] ` `=` `1` ` ` ` ` `return` `res ` ` ` `# Driver Code ` `if` `__name__ ` `=` `=` `"__main__"` `: ` ` ` `string ` `=` `"0102010"` ` ` `print` `(getSubstringWithEqual012(string)) ` ` ` `# This code is conributed by ` `# sanjeev2552 ` |

*chevron_right*

*filter_none*

Output:

2

This article is contributed by **Utkarsh Trivedi**. If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please write comments if you find anything incorrect, or you want to share more information about the topic discussed above.

Don’t stop now and take your learning to the next level. Learn all the important concepts of Data Structures and Algorithms with the help of the most trusted course: **DSA Self Paced**. Become industry ready at a student-friendly price.

## Recommended Posts:

- Count of substrings whose Decimal equivalent is greater than or equal to K
- Split the binary string into substrings with equal number of 0s and 1s
- Number of substrings with count of each character as k
- Count number of distinct substrings of a given length
- Count number of substrings with numeric value greater than X
- Count number of substrings with exactly k distinct characters
- Count number of substrings of a string consisting of same characters
- Given a binary string, count number of substrings that start and end with 1.
- Count the number of vowels occurring in all the substrings of given string
- Count number of triplets with product equal to given number
- Count triplets such that sum of any two number is equal to third | Set 2
- Count subarrays with equal number of 1's and 0's
- Count number of equal pairs in a string
- Count the number of strings in an array whose distinct characters are less than equal to M
- Split a given string into substrings of length K with equal sum of ASCII values
- Count number of binary strings such that there is no substring of length greater than or equal to 3 with all 1's
- Count substrings with same first and last characters
- Count all substrings having character K
- Count substrings that contain all vowels | SET 2
- Count of substrings consisting only of vowels