# 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)