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

Difficulty Level : Hard
Last Updated : 06 Aug, 2021

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:

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready. To complete your preparation from learning a language to DS Algo and many more, please refer Complete Interview Preparation Course.

In case you wish to attend live classes with experts, please refer DSA Live Classes for Working Professionals and Competitive Programming Live for Students.

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²) checking each substring for count takes O(n) times, So total time taken to solve this problem is O(n³) 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;
}

Java

// Java program to find substring with equal
// number of 0's, 1's and 2's
import java.io.*;
import java.util.*;
 
class GFG {
 
    // Method to count number of substring which
    // has equal 0, 1 and 2
    private static int getSubstringWithEqual012(String str)
    {
 
        // map to store, how many times a difference
        // pair has occurred previously (key = diff1 *
        // diff2)
        HashMap<String, Integer> map = new HashMap<>();
        map.put("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;
 
        int ans = 0;
 
        //  looping into string
        for (int i = 0; i < str.length(); i++) {
            // increasing the count of current character
            if (str.charAt(i) == '0')
                zc++;
            else if (str.charAt(i) == '1')
                oc++;
            else
                tc++;
 
            // making key of differences (z[i] - o[i],
            // z[i] - t[i])
            String key = (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
            ans += map.getOrDefault(key, 0);
 
            map.put(key, map.getOrDefault(key, 0) + 1);
        }
        // increasing the count of current difference
        // pair by 1
        return ans;
    }
 
    // Driver Code
    public static void main(String[] args)
    {
        // Input
        String str = "0102010";
        System.out.println(getSubstringWithEqual012(str));
    }
}

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 contributed by
# sanjeev2552

Output:

This article is contributed by Utkarsh Trivedi. If you like GeeksforGeeks and would like to contribute, you can also write an article using write.geeksforgeeks.org or mail your article to review-team@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.

My Personal Notes

Related Articles

Related Articles

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

C++

Java

Python3