Solution to 'minimum number of of edits needed to assert r'A+B+'

mwenja-solution.py

# Solve: minimum number of 'edits' needed to assert r'A+B+' in <string>

def main(string):
    # count of total A's, B's
    totals = {"A": 0, "B": 0}
    # count of  A's, B's seen (in second loop) (does not include edits)
    total_seen = {"A": 0, "B": 0}
    # solution/answer
    solution = 0

    # O(n)
    for i in string:
        totals[i] += 1

    # track current char
    char = string[0]
    # count similar char
    total_char = 0

    # O(n)
    for i in string:
        if i == char:
            # count similar char
            total_char += 1
            continue

        if char == "A":
            # look behind, if any B exists, means we have to go!
            if total_seen["B"] > 0:
                # remove A's
                solution += total_char
            else:
                # count A's seen
                total_seen["A"] += total_char

        if char == "B":
            # look ahead, if any A exists larger then our B's, we have to go!
            if total_char < (totals["A"] - total_seen["A"]):
                # remove B's
                solution += total_char
            else:
                # count B's seen
                total_seen["B"] += total_char

        # reset char, total_char
        char = i
        total_char = 1
    
    # O(1)
    # final check (last char)
    if char == "A" and total_seen["B"] > 0:
        # remove A's
        solution += total_char
    
    # should not happen ;)
    if solution > max(totals["A"], totals["B"]):
        raise Exception(
            f"solution exceeds worst case: total_A: {totals['A']}, total_B: "
            f"{totals['B']}, solution: {solution}"
        )
    
    # O(n + n + 1) == O(n2 + 1) == O(n2)?
    return solution
    
    
if __name__ == "__main__":
    strings = [
        'A',
        'B',
        'AA',
        'AB',
        'BA',
        'BB',
        'AAA',
        'AAB',
        'ABB',
        'ABA',
        'BAA',
        'BBB',
        'BBA',
        'BAB',
        'BBB',
        'BAAABAB',
        'ABBAAAABBBBAAB',
        'BAAABABABABAABABABABABABABABBBABABAAAABBBAAAABB',
    ]
    
    for s in strings:
        print(s, main(s))