Problem Statement in English

You’re given a string s and two integers x and y. You can remove substrings “ab” or “ba” from s, earning x points for each “ab” removed and y points for each “ba” removed. Your task is to maximize the score by removing these substrings optimally.

Approach

We need to iterate through the string and use a stack to keep track of characters. When we encounter “ab” or “ba”, we pop them from the stack and add the corresponding score to our total score. We will prioritize removing the substring that gives us a higher score first.

The stack just makes it easier to handle the removal of characters as we can easily check the last two characters in the stack to see if they form “ab” or “ba”.

Solution in Python

class Solution:
    def maximumGain(self, s: str, x: int, y: int) -> int:
        stack = []
        score = 0

        if x > y:
            higher_list = ['a', 'b']
            lower_list = ['b', 'a']
            higher_score = x
            lower_score = y
        else:
            higher_list = ['b', 'a']
            lower_list = ['a', 'b']
            higher_score = y
            lower_score = x

        for c in s:
            stack.append(c)

            if stack[-2:] == higher_list:
                stack.pop()
                stack.pop()
                score += higher_score
        
        nstack = []
        for c in stack:
            nstack.append(c)

            if nstack[-2:] == lower_list:
                nstack.pop()
                nstack.pop()
                score += lower_score

        return score

Complexity

  • Time: $O(n)$
    Since we traverse the string once with two pointers, the time complexity is linear.

  • Space: $O(n)$
    The space complexity is linear because we use a stack to store the characters.

Mistakes I Made

I didn’t think it was necessary to use the stack for the second pass, but it simplifies the logic and makes it easier to handle the removal of characters.

And we are done.