Count Substrings with equal number of 0s, 1s and 2s
Last Updated :
24 Apr, 2025
Given a string that consists of only 0s, 1s and 2s, count the number of substrings that have an 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
Explanation: Substring str[0, 2] = "102", str[1, 3] = "021", str[2, 4] = "210", str[2, 7] = "210021", and str[5, 7] = "021"
each have an equal number of 0s, 1s, and 2s.
[Naive Approach] Using Nested Loops – O(n^3) Time and O(1) Space
Iterate through all substrings of string using nested loops and check whether they contain equal 0,1 and 2 or not.
C++
#include <bits/stdc++.h>
using namespace std;
int countSubstrs(string s) {
int n = s.size();
int ct = 0;
for (int i = 0; i < n; ++i) {
for (int j = i + 1; j <= n; ++j) {
// Extract substring from i to j
string substr = s.substr(i, j - i);
int ct0 = 0, ct1 = 0, ct2 = 0;
for (char c : substr) {
if (c == '0') ct0++;
if (c == '1') ct1++;
if (c == '2') ct2++;
}
if (ct0 == ct1 && ct1 == ct2) {
ct++;
}
}
}
return ct;
}
int main() {
string s = "0102010";
cout << countSubstrs(s) << endl;
return 0;
}
Java
import java.util.*;
public class Main {
public static int countSubstrs(String s) {
int n = s.length();
int ct = 0;
for (int i = 0; i < n; ++i) {
for (int j = i + 1; j <= n; ++j) {
// Extract substring from i to j
String substr = s.substring(i, j);
int ct0 = 0, ct1 = 0, ct2 = 0;
for (char c : substr.toCharArray()) {
if (c == '0') ct0++;
if (c == '1') ct1++;
if (c == '2') ct2++;
}
if (ct0 == ct1 && ct1 == ct2) {
ct++;
}
}
}
return ct;
}
public static void main(String[] args) {
String s = "0102010";
System.out.println(countSubstrs(s));
}
}
Python
def countSubstrs(s):
n = len(s)
ct = 0
for i in range(n):
for j in range(i + 1, n + 1):
# Extract substring from i to j
substr = s[i:j]
ct0, ct1, ct2 = 0, 0, 0
for c in substr:
if c == '0': ct0 += 1
if c == '1': ct1 += 1
if c == '2': ct2 += 1
if ct0 == ct1 and ct1 == ct2:
ct += 1
return ct
s = "0102010"
print(countSubstrs(s))
C#
using System;
class Program {
public static int countSubstrs(string s) {
int n = s.Length;
int ct = 0;
for (int i = 0; i < n; ++i) {
for (int j = i + 1; j <= n; ++j) {
// Extract substring from i to j
string substr = s.Substring(i, j - i);
int ct0 = 0, ct1 = 0, ct2 = 0;
foreach (char c in substr) {
if (c == '0') ct0++;
if (c == '1') ct1++;
if (c == '2') ct2++;
}
if (ct0 == ct1 && ct1 == ct2) {
ct++;
}
}
}
return ct;
}
static void Main() {
string s = "0102010";
Console.WriteLine(countSubstrs(s));
}
}
JavaScript
function countSubstrs(s) {
let n = s.length;
let ct = 0;
for (let i = 0; i < n; ++i) {
for (let j = i + 1; j <= n; ++j) {
// Extract substring from i to j
let substr = s.substring(i, j);
let ct0 = 0, ct1 = 0, ct2 = 0;
for (let c of substr) {
if (c === '0') ct0++;
if (c === '1') ct1++;
if (c === '2') ct2++;
}
if (ct0 === ct1 && ct1 === ct2) {
ct++;
}
}
}
return ct;
}
let s = "0102010";
console.log(countSubstrs(s));
[Expected Approach] Using Hashing – O(n) Time and O(n) Space
Traverse through the string and keep track of counts of 0, 1, and 2 and make a difference pair of (zeroes - ones, zeroes - twos) and increase the answer count if this difference pair is seen before and at every index increase the count of this difference pair in the map
Follow the given steps to solve the problem:
- Declare a map to store the difference pair and three variables to store the count of 0's, 1's and 2's
- Traverse the string and keep track of the count of 0's, 1's, and 2's
- At each index make a difference pair of (zeroes - ones, zeroes - twos)
- Using the map check if this pair is seen before, if it is so then increase the result count
- Then, increase the count of this pair in the map
- Return the result
C++
#include <bits/stdc++.h>
using namespace std;
int countSubstrs(string s) {
int n = s.size(), cnt = 0;
map<pair<int, int>, int> mp;
int z = 0, o = 0, t = 0;
// Initialize the map with the base case (0, 0)
mp[{0, 0}] = 1;
for (int i = 0; i < n; ++i) {
if (s[i] == '0') z++;
if (s[i] == '1') o++;
if (s[i] == '2') t++;
// difference pair (z - o, z - t)
pair<int, int> diff = {z - o, z - t};
// If this pair has been seen before, add its frequency to the count
cnt += mp[diff];
mp[diff]++;
}
return cnt;
}
int main() {
string s = "102100211";
cout << countSubstrs(s) << endl;
return 0;
}
Java
import java.util.*;
public class Main {
public static int countSubstrs(String s) {
int n = s.length(), cnt = 0;
Map<String, Integer> mp = new HashMap<>();
int z = 0, o = 0, t = 0;
// Initialize the map with the base case (0, 0)
mp.put("0,0", 1);
for (int i = 0; i < n; i++) {
if (s.charAt(i) == '0') z++;
if (s.charAt(i) == '1') o++;
if (s.charAt(i) == '2') t++;
// Create the difference pair (z - o, z - t)
String diff = (z - o) + "," + (z - t);
// If this pair has been seen before, add its frequency to the count
cnt += mp.getOrDefault(diff, 0);
mp.put(diff, mp.getOrDefault(diff, 0) + 1);
}
return cnt;
}
public static void main(String[] args) {
String s = "102100211";
System.out.println(countSubstrs(s));
}
}
Python
from collections import defaultdict
def countSubstrs(s):
n = len(s)
cnt = 0
mp = defaultdict(int)
z, o, t = 0, 0, 0
# Initialize the map with the base case (0, 0)
mp[(0, 0)] = 1
for i in range(n):
if s[i] == '0': z += 1
if s[i] == '1': o += 1
if s[i] == '2': t += 1
# Create the difference pair (z - o, z - t)
diff = (z - o, z - t)
# If this pair has been seen before, add its frequency to the count
cnt += mp[diff]
mp[diff] += 1
return cnt
s = "102100211"
print(countSubstrs(s))
C#
using System;
using System.Collections.Generic;
class Program {
public static int countSubstrs(string s) {
int n = s.Length, cnt = 0;
var mp = new Dictionary<string, int>();
int z = 0, o = 0, t = 0;
// Initialize the map with the base case (0, 0)
mp["0,0"] = 1;
for (int i = 0; i < n; i++) {
if (s[i] == '0') z++;
if (s[i] == '1') o++;
if (s[i] == '2') t++;
string diff = (z - o) + "," + (z - t);
// If this pair has been seen before, add its frequency to the count
if (mp.ContainsKey(diff)) {
cnt += mp[diff];
}
if (mp.ContainsKey(diff)) {
mp[diff]++;
} else {
mp[diff] = 1;
}
}
return cnt;
}
static void Main() {
string s = "102100211";
Console.WriteLine(countSubstrs(s));
}
}
JavaScript
function countSubstrs(s) {
let n = s.length;
let cnt = 0;
let mp = {};
let z = 0, o = 0, t = 0;
// Initialize the map with the base case (0, 0)
mp["0,0"] = 1;
for (let i = 0; i < n; i++) {
if (s[i] === '0') z++;
if (s[i] === '1') o++;
if (s[i] === '2') t++;
// Create the difference pair (z - o, z - t)
let diff = (z - o) + "," + (z - t);
// If this pair has been seen before, add its frequency to the count
if (mp[diff]) {
cnt += mp[diff];
}
if (mp[diff]) {
mp[diff]++;
} else {
mp[diff] = 1;
}
}
return cnt;
}
let s = "102100211";
console.log(countSubstrs(s));
The above C++ code is O(n Log n) as we us a map (Which is self balancing tree based) rather than unordered_map. C++ does not support unordered_map with pair of ints as key. However, we can use long long as a hack in C++ to achieve O(n) Time. We mainly store both values in a single long long variable.
C++
#include <bits/stdc++.h>
using namespace std;
int countEqualSubstrings(string s) {
int n = s.size(), cnt = 0, z = 0, o = 0, t = 0;
unordered_map<long long, int> mp;
// Initialize the map with the base case (0,0)
mp[0] = 1;
for (char c : s) {
if (c == '0') z++;
if (c == '1') o++;
if (c == '2') t++;
// Create a unique hash for the difference pair (z - o, z - t)
long long diff = (long long)(z - o) * 1000000007 + (z - t);
// If this pair has been seen before, add
// its frequency to the count
cnt += mp[diff];
mp[diff]++;
}
return cnt;
}
int main() {
string s = "102100211";
cout << countEqualSubstrings(s) << endl;
}
Similar Reads
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 an equal number of 0's and 1's. Examples: Input: arr[] = {1, 0, 0, 1, 0, 1, 1}Output: 8Explanation: The index range for the 8 sub-arrays are: (0, 1), (2, 3), (0, 3), (3, 4), (4, 5)(2, 5), (0, 5), (1,
14 min read
Count Substrings with number of 0s and 1s in ratio of X : Y Given a binary string S, the task is to count the substrings having the number of 0s and 1s in the ratio of X : Y Examples: Input: S = "010011010100", X = 3, Y = 2Output: 5Explanation: The index range for the 5 substrings are: (0, 4), (2, 6), (6, 10), (7, 11), (2, 11) Input: S = "10010101", X = 1, Y
7 min read
Longest balanced binary substring with equal count of 1s and 0s Given a binary string str[] of size N. The task is to find the longest balanced substring. A substring is balanced if it contains an equal number of 0 and 1. Examples: Input: str = "110101010"Output: 10101010Explanation: The formed substring contain equal count of 1 and 0 i.e, count of 1 and 0 is sa
8 min read
Counting Good Substrings with 0s, 1s, and ?s Given a string S consisting of characters 0,1,? A string is called good if all the adjacent characters in it are different i.e. of the form 0101010... or 1010101... You can replace ? with 0 or 1. Your task is to determine the number of good contiguous substrings of the string S. Constraint: |S|<2
7 min read
Split the binary string into substrings with equal number of 0s and 1s Given a binary string str of length N, the task is to find the maximum count of consecutive substrings str can be divided into such that all the substrings are balanced i.e. they have equal number of 0s and 1s. If it is not possible to split str satisfying the conditions then print -1.Example: Input
8 min read