python-codewars-solutions/simplify-poly/solution.py

# https://www.codewars.com/kata/55f89832ac9a66518f000118/

import re
from dataclasses import dataclass


@dataclass
class Term:
    coef: int
    variables: str

    def __repr__(self):
        return f"Term(coef={self.coef}, vars={self.variables})"

    def __str__(self):
        a = str(self.coef) if self.coef != 1 else ""
        a = "-" if a == "-1" else a
        b = "".join(self.variables)
        return f"{a}{b}"

    def __add__(self, other):
        if self.variables == other.variables:
            r = Term(self.coef + other.coef, self.variables)
            return r
        raise ArithmeticError("Variables do not match")

    def __lt__(self, other):
        if len(self.variables) == len(other.variables):
            return self.variables < other.variables

        return len(self.variables) < len(other.variables)

    def __eq__(self, other):
        return self.coef == other.coef and self.variables == other.variables


def parse_term(inp: str) -> Term:
    sign, coef, variables = "", [], []
    for n in inp:
        if n == "-":
            sign = "-"
        elif n.isnumeric():
            coef.append(n)
        elif n.isalpha():
            variables.append(n)

    c = -1 if sign == "-" else 1
    c *= int("".join(coef)) if coef else 1
    return Term(c, sorted(variables))


def simplify_terms(terms: [Term]) -> [Term]:
    simplified_terms = []
    for term in terms:
        if not simplified_terms:
            simplified_terms.append(term)
            continue

        previous = simplified_terms.pop()
        try:
            res = previous + term
            if res.coef != 0:
                simplified_terms.append(res)
        except ArithmeticError:  # terms not compatible to be added
            simplified_terms.append(previous)
            simplified_terms.append(term)
    return simplified_terms


def poly_to_str(terms: [Term]) -> str:
    res = ""
    for term in terms:
        s = str(term)
        if res and not s.startswith("-"):
            res += "+"
            res += s
        else:
            res += s
    return res


def simplify(poly: str):
    pattern = r"([+-]?\d*\w+)"
    terms = sorted(parse_term(t) for t in re.findall(pattern, poly))
    simplified_terms = simplify_terms(terms)
    return poly_to_str(simplified_terms)