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

The problem “Count Substrings with equal number of 0s, 1s and 2s” states that you are given a string that has 0, 1, and 2 only. The problem statement asks to find out the number of substrings that contain equal no of 0, 1, and 2 only.

Table of Contents

Example

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

str= “01200”

Explanation

Two substrings, that have an equal number of 0, 1, and 2 are (012) and (120).

str= “10201012”

Explanation

Four substrings, that have an equal number of 0, 1, and 2 are (102) and (201), (012) and (201012).

Algorithm

Declare a map.
Initialize a pair to 0 and put it into the map with 1.
Set zerocount, onecount, twocount, and output to 0.
Loop from i=0 to i<n.
1. Check for each character of the current substring whether it is 0, 1, and 2. Then update the count accordingly.
2. Calculate the differences of pair.
3. Check if the difference is not present in the map, then, add 0 to the output.
4. Else, add temp’s value of map into the output.
5. Add temp to the map.
Return output.

Explanation

We are given a string that has 0, 1, and 2 only. Our task is to find out the total number of the substrings that have equal no of 0, 1, and 2. For this, we are going to use Hashing. Initialize a pair with (0, 0) as key and its value as 1 into the map, by default. Calculate the difference between zerocount and onecount, and zerocount and twocount. We will be storing the value in a pair and that pair into the map. If the difference of a pair already exists in the map, then simply get/retrieve the value of the current pair from the map. Then add that to the output. If the difference is not already present in the map. Then add 0 to the output. We also need to insert the difference pair into the map and increase its frequency if it already exists in the map. Else store a new entry for the difference pair with 1 as a value into the map.

Let us consider an example for it:

Example

String = “01200”

We have already initialized pair to 0 and insert pair as key and 1 as value into the map. We will retrieve the first character like 0. So we will calculate the difference pair and get (1,1). Also, this is because we have increased the count of zero’s. So we will calculate the difference and get that result. After getting the value of the map’s current value and adding 0 to the output, because this pair is new. We will just insert the pair into the map. The current output is 0.

We will go for the next character that is 1. Now we will increase the count of one’s. Then we will get the differences as 0,1. Since this pair is also new, so we will add that to the map and the output still remains the same.

Then we get 2 as an input we will get the pair as 0, 0 because all of the digit’s count is now 1. We will store that too in the map and this time update output because 0,0 we have already initialized in a map so the output is 1 now.

Next, we get like 0, now the zerocount is two. We get the 1,1 as this is already in the map. Then we will update the output and insert it in the map with the value 2.

At this point we have found our all our possible substrings, this is the way we get equal no of 0’s, 1’s and 2’s.

Code

C++ code to count substrings with equal number of 0s, 1s and 2s

#include<bits/stdc++.h>
using namespace std;

struct hash_pair { 
    template <class T1, class T2> 
    size_t operator()(const pair<T1, T2>& p) const
    { 
        auto hash1 = hash<T1>{}(p.first); 
        auto hash2 = hash<T2>{}(p.second); 
        return hash1 ^ hash2; 
    } 
}; 

int getSubString(string str)
{
    int n = str.length();
    unordered_map< pair<int, int>, int, hash_pair > MAP;
    MAP[make_pair(0, 0)] = 1;
    int zerocount = 0, onecount = 0, twocount = 0;
    int output = 0;
    for (int i = 0; i < n; ++i)
    {
        if (str[i] == '0')
            zerocount++;
        else if (str[i] == '1')
            onecount++;
        else
            twocount++;
        pair<int, int> x = make_pair(zerocount - onecount,zerocount - twocount);
        output = output + MAP[x];
        MAP[x]++;
    }
    return output;
}
int main()
{
    string str = "10201012";
    cout << getSubString(str) << endl;
    return 0;
}

Java code to count substrings with equal number of 0s, 1s and 2s

import java.util.HashMap;
class Pair
{
    int first, second;
    Pair(int x, int y)
    {
        this.first=x;
        this.second=y;
    }
}
class SubstringCount
{
    public static int getSubString012(String str1)
    {
        int n = str1.length();

        HashMap< String, Integer > MAP=new HashMap<>();



        MAP.put("0:0", 1);

        int zerocount = 0, onecount = 0, twocount = 0;

        int output = 0;
        for (int i = 0; i < n; ++i)
        {
            if (str1.charAt(i)=='0')
                zerocount++;
            else if (str1.charAt(i) == '1')
                onecount++;
            else
                twocount++;

            Pair pair = new Pair(zerocount - onecount, zerocount - twocount);

            String str=pair.first+":"+pair.second;

            if(!MAP.containsKey(str))
                output = output + 0;
            else
                output = output + MAP.get(str);

            if(MAP.containsKey(str))
                MAP.put(str, MAP.get(str)+1);
            else
                MAP.put(str, 1);
        }

        return output;
    }
    public static void main(String [] args)
    {
        String str = "10201012";
        System.out.println(getSubString012(str));
    }
}

Complexity Analysis

Time Complexity

O(n) where “n” is the number of elements in the array. Since we have used HashMap, all of the insertion, deletion, and searching require only O(1) time per operation.

Space Complexity

O(n) where “n” is the number of elements in the array.