Designing an Algorithm for Minimum Substring Removal to Balance a String of ‘a’ and ‘b’

Before writing pseudocode, we must first analyze the problem carefully and identify an efficient strategy that satisfies the input constraints. The following steps outline the reasoning process used to simplify the problem and construct the final algorithm.

Step 1: Identify the problem statement

We are given a string consisting only of the characters ‘a’ and ‘b’. The goal is to remove as few consecutive characters as possible so that the remaining string contains an equal number of ‘a’ and ‘b’. Since only consecutive characters may be removed, the problem reduces to finding the shortest contiguous substring whose removal makes the string balanced.

Understanding the problem statement is often the hardest part of programming assignments. In fact, learning how to understand your programming homework is a key skill before writing any algorithm.

Step 2: Analyze Input Constraints

The maximum length of a string in a single test case is 2 * 10 5. Additionally, the total length across all
test cases do not exceed 2 × 10⁵. These constraints are significant. A quadratic solution O(n2). It would be too slow for n up to 200,000. Therefore, the algorithm must operate in linear time O(n) or at most O(n log n).

Step 3:  Reduce the problem

In the case of an unbalanced string, the count of one letter will always be greater than that of the other. Consider the string “aabaabababa” as an example. Here, the count of a is 7 and the count of b is 4. If we think greedily, which letter should we remove? We should focus on removing an a. If we remove a b, we would then need to remove an additional a to compensate. Now, let us denote the number of occurrences of the letter a by A and the number of occurrences of the letter b by B. Let their difference be X, where X = A – B.

Assume we have found the minimum substring that must be removed to make the string balanced. What would this substring look like? Since we need to remove X occurrences of the letter a, the substring must contain at least X a’s. If these X a’s are not contiguous, the substring will also include some b’s. For each such b, we must remove an additional a to offset it. Thus, the substring will contain X a’s plus some equal number of additional a’s and b’s, which effectively cancel each other out, leaving a net excess of X a’s. Based on this observation, the problem reduces to finding a minimum-length substring in which the difference between the number of occurrences of a and b is equal to X.

Step 4: Constructing the algorithm

We need an efficient way to compute the difference for any substring. To do this, we can use a prefix difference array. Each element in this array represents the difference between the total number of occurrences of a and b up to that index. Then, the difference for a substring from index l to r can be computed as difference[r] – difference[l − 1]. 

Suppose the required substring starts at index l and ends at index r. Then its difference value is difference[r] – difference[l − 1], which must be equal to X. To find such a substring, we can iterate through each index and treat it as r. For each r, we ask: what value must be subtracted from difference[r] so that the result equals X? Let this value be Y. Then difference[r] – Y = X, so Y = difference[r] – X.

Thus, for each r, we only need to find the nearest index till r where the prefix difference equals Y. We can maintain a hashmap in which the key is the difference value and the value is the index. If a difference value already exists in the hashmap, we replace it, since the newer index will always be closer to r.

Let’s write the pseudocode for this idea. 

SET array difference[0 to n - 1]
Compute prefix difference array

SET answer = +INFINITY
Create empty hash table H

FOR r = 0 TO n - 1 DO
    SET diff = difference[r]
    SET H[diff] = r
    SET Y = diff - X

    IF Y exists in H THEN
        SET l = H[Y]
        SET answer = MIN(answer, r - l)
    ELSE IF Y = 0 THEN
        SET answer = MIN(answer, r + 1)
    END IF
END FOR

Note the edge case: if Y is not present in the hash table and Y = 0, then it implies that the entire prefix from index 0 to r must be removed. This algorithm runs in O(n log n) time due to the use of a hash map. It can be optimized to O(n) if we instead use an array, where the index of the array represents Y and the value stored at that index represents the position. Since Y can be negative, we can offset it by N and increase the size of the array to 2N + 1.This shows how focusing on the input constraints helps in designing a more efficient solution.

Let us combine these ideas and construct a generic pseudocode. Instead of explicitly using specific characters, we define maxChar as the character with the greater number of occurrences and minChar as the character with the smaller number of occurrences.

Step 5: Pseudocode

SOLVE-TESTCASE(S, n)
    SET countA = 0
    SET countB = 0

    // Find the frequency of letter 'a' and letter 'b'
    FOR i = 0 TO n - 1 DO
        IF S[i] = 'a' THEN
            SET countA = countA + 1
        ELSE
            SET countB = countB + 1
        END IF
    END FOR

    // Define maxChar as the letter with the highest frequency
    // and minChar as the letter with the lowest frequency

    IF countA > countB THEN
        SET maxChar = 'a'
        SET minChar = 'b'
        SET X = countA - countB
    ELSE
        SET maxChar = 'b'
        SET minChar = 'a'
        SET X = countB - countA
    END IF

    SET answer = +INFINITY
    Create empty hash table H
    SET minCharCount = 0
    SET maxCharCount = 0

    FOR r = 0 TO n - 1 DO
        // Get the running frequency and difference at index r
        IF S[r] = minChar THEN
            SET minCharCount = minCharCount + 1
        ELSE
            SET maxCharCount = maxCharCount + 1
        END IF

        SET difference = maxCharCount - minCharCount
        SET H[difference] = r
        SET Y = difference - X
   
        IF Y EXISTS in H THEN
            SET l = H[Y]
            SET answer = MIN(answer, r - l)
        ELSE IF Y = 0 THEN
            SET answer = MIN(answer, r + 1)
        END IF
    END FOR

    RETURN answer

Here is the implementation of this pseudocode in C++: https://pastebin.com/jU84WzER