Leetcode 640 Solution

class Solution:
    def split_equation(self, equation):
        i = 0
        ans = []
        while i < len(equation):
            j = i
            while j < len(equation) - 1 and equation[j + 1] not in ['-', '+']:
                j += 1
            ans.append(equation[i:j + 1])
            i = j + 1
        return ans

    def coeff(self, v):
        if v == 'x' or v == '+x':
            return 1
        elif v == '-x':
            return -1
        else:
            return int(v[:-1].lstrip('+'))

    def solveEquation(self, equation: str) -> str:
        left_part, right_part = equation.split('=')
        left_tokens = self.split_equation(left_part)
        right_tokens = self.split_equation(right_part)

        x_coefficient = 0
        coefficient = 0
        for token in left_tokens:
            if token[-1] == 'x':
                x_coefficient += self.coeff(token)
            else:
                coefficient += int(token)

        for token in right_tokens:
            if token[-1] == 'x':
                x_coefficient -= self.coeff(token)
            else:
                coefficient -= int(token)

        if x_coefficient == 0:
            return "No solution" if coefficient != 0 else "Infinite solutions"
        else:
            return "x={}".format(-coefficient // x_coefficient)