Problem: Equal 0, 1 and 2

#Difficulty: Easy

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

Implementation: Python2.7

"""
    @author: loctv
    @time: January 2, 2017
"""
def all_sub_strings(string):
    """
        generate all consecutive substrings of a given string
    """
    length = len(string)
    if length == 3:
        yield string
    else:
        for i in range(length):
            for j in range(i, length+1):
                yield string[i:j]

def is_equal(string):
    """
        check if inside a string, frequency of
        1, 0, and 2 must be equal
        and their frequency must be bigger than zero
    """
    frequency = {'0':0, '1':0, '2':0}
    for _ in string:
        frequency[_] += 1
    return  frequency['0'] == frequency['1'] == frequency['2'] and frequency['0'] > 0

if __name__ == '__main__':
    TEST_CASE = input()
    for _ in range(TEST_CASE):
        string = raw_input().strip()
        total = 0
        for sub_string in all_sub_strings(string):
            #print sub_string
            total += 1 if is_equal(sub_string) else 0
        print total

    
Advertisements