Equal 0, 1 and 2

Given a string which consists of only 0, 1 or 2s, count the number of substring which have 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

Input:

The first line of input contains a single integer T denoting the number of test cases. Then T test cases follow. Each test case consists of one line. The line contains a string of 0's, 1's and 2's.

Output:

Corresponding to each test case, in a new line, print the count all possible substrings that have same number of 0s, 1s and 2s.

Constraints:

1 ≤ T ≤ 100 1 ≤ string length ≤ 1000 Example:

Input 2 0102010 102100211

Output 2 5

class Solution {
    public static int numberOfSubStrings(String str) {
        // Frequency array to store frequency of the 3 digits
        int[] freq = new int[3];
        // Map of pair difference string to frequency
        Map<String, Integer> map = new HashMap<>();
        // Default difference in pairs = 0
        map.put("00", 1);
        int count = 0;
        for (char c : str.toCharArray()) {
            int digit = (int) (c - '0');
            freq[digit]++;
            // The key will be of this form
            String key = Integer.toString(freq[1] - freq[0]) + Integer.toString(freq[2] - freq[1]);
            if (map.containsKey(key))
                count += map.get(key);
            map.put(key, map.getOrDefault(key, 0) + 1);
        }
        return count;
    }
}