Program to find out the length of the substring where two times the number of zeroes in substring is lesser than or equal to three times the number of ones in the substring in Python

Sometimes we need to find the maximum length substring where the number of zeros and ones satisfy a specific condition. This problem involves finding a substring in a repeated string where 2 × (zeros) ? 3 × (ones).

Problem Understanding

Given a string and integer k, we repeat the string k times to create a new string. We need to find the longest substring where the condition 2 * zeros ? 3 * ones holds true.

Example

If k = 2 and input_str = '0101011', the repeated string becomes '01010110101011' (length 14). This entire string satisfies our condition: 6 zeros and 8 ones, where 2 × 6 ? 3 × 8 (12 ? 24).

Algorithm Steps

The solution uses prefix sums and binary search to efficiently find the maximum valid substring ?

• Create prefix arrays to track the balance between zeros and ones
• Transform the condition 2 × zeros ? 3 × ones into a cumulative sum problem
• Use binary search to find valid substring endpoints efficiently
• Handle the repeated string pattern to cover all possible substrings

Implementation

from bisect import bisect_left

def solve(k, input_str):
    str_len = len(input_str)
    list_a = [0] * (str_len + 1)
    list_b = [0] * (str_len + 1)
    list_b[0] = (0, 0)
    
    # Build prefix sum array
    for i in range(str_len):
        list_a[i + 1] = list_a[i] - 3 * (input_str[i] == '1') + 2 * (input_str[i] == '0')
        list_b[i + 1] = (list_a[i + 1], i + 1)
    
    # Sort by prefix sum values
    list_b.sort()
    
    # Build maximum position array
    temp_list = [0] * (str_len + 1)
    temp_list[0] = list_b[0][1]
    for i in range(str_len):
        temp_list[i + 1] = max(temp_list[i], list_b[i + 1][1])
    
    res = 0
    
    # Find maximum valid substring length
    for i in range(str_len):
        tmp = list_b[0][0] - list_a[i]
        
        if list_a[str_len] <= 0:
            a = k - 1
            if tmp + list_a[str_len] * a > 0:
                continue
        elif tmp > 0:
            continue
        else:
            a = min(k - 1, -tmp // list_a[str_len])
        
        v = a * list_a[str_len] - list_a[i]
        b = bisect_left(list_b, (-v + 1, 0)) - 1
        res = max(res, temp_list[b] - i + a * str_len)
    
    return res

# Test the function
print(solve(2, '0101011'))

14

How It Works

The algorithm transforms the inequality 2 × zeros ? 3 × ones into a prefix sum problem. For each character, we add +2 for '0' and -3 for '1'. A valid substring exists when the prefix sum difference is non-positive.

Key Points

• The prefix sum approach converts the counting problem into range queries
• Binary search optimizes finding valid endpoints in the sorted array
• The algorithm handles multiple repetitions of the string efficiently
• Time complexity is O(n log n) where n is the string length

Conclusion

This solution efficiently finds the maximum substring length using prefix sums and binary search. The key insight is transforming the counting constraint into a cumulative sum problem that can be solved with standard algorithms.