Open In App

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

Last Updated : 24 Apr, 2025
Comments
Improve
Suggest changes
Like Article
Like
Report

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));

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
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)); 

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.

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;
}

Output
5

Next Article
Article Tags :
Practice Tags :

Similar Reads