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 2Input: 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.
Try It Yourself
Table of Content
[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.
#include <bits/stdc++.h>
using namespace std;
int equal012(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 << equal012(s) << endl;
return 0;
}
import java.util.*;
public class Main {
public static int equal012(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(equal012(s));
}
}
def equal012(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(equal012(s))
using System;
class Program {
public static int equal012(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(equal012(s));
}
}
function equal012(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(equal012(s));
Output
2
[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
#include <bits/stdc++.h>
using namespace std;
int equal012(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 << equal012(s) << endl;
return 0;
}
import java.util.*;
public class Main {
public static int equal012(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(equal012(s));
}
}
from collections import defaultdict
def equal012(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(equal012(s))
using System;
using System.Collections.Generic;
class Program {
public static int equal012(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(equal012(s));
}
}
function equal012(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(equal012(s));
Output
5
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.
#include <bits/stdc++.h>
using namespace std;
int equal012(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 << equal012(s) << endl;
}
Output
5