simplifying multilinear polynomials

simplify-poly/solution.py

@@ -3,14 +3,15 @@
 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
@@ -18,47 +19,37 @@ class Term:
         return f"{a}{b}"
 
     def __add__(self, other):
-        print(f"add {self} {other}")
         if self.variables == other.variables:
-            r = Term(self.coef+other.coef, self.variables)
-            print(r)
+            r = Term(self.coef + other.coef, self.variables)
             return r
-        else:
-            print(self.variables, other.variables)
         raise ArithmeticError("Variables do not match")
 
     def __lt__(self, other):
-        return (len(self.variables) < len(other.variables)) and (self.variables < other.variables)
-    
+        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 = '', [], []
+    sign, coef, variables = "", [], []
     for n in inp:
-        if n == '-':
-            sign = '-'
+        if n == "-":
+            sign = "-"
         elif n.isnumeric():
             coef.append(n)
         elif n.isalpha():
-            variables.append(n) 
-    
-    c = -1 if sign == '-' else 1
+            variables.append(n)
+
+    c = -1 if sign == "-" else 1
     c *= int("".join(coef)) if coef else 1
     return Term(c, sorted(variables))
 
 
-def simplify(poly):
-    print(poly)
-    
-    # parse
-    pattern = r"([+-]?\d*\w+)"
-    terms = sorted(parse_term(t) for t in re.findall(pattern, poly))
-    
-    print(terms)
-    
-    # simplify
+def simplify_terms(terms: [Term]) -> [Term]:
     simplified_terms = []
     for term in terms:
         if not simplified_terms:
@@ -70,14 +61,15 @@ def simplify(poly):
             res = previous + term
             if res.coef != 0:
                 simplified_terms.append(res)
-        except ArithmeticError as e:
-            print("no match")
+        except ArithmeticError:  # terms not compatible to be added
             simplified_terms.append(previous)
             simplified_terms.append(term)
-    
-    # build output str
+    return simplified_terms
+
+
+def poly_to_str(terms: [Term]) -> str:
     res = ""
-    for term in simplified_terms:
+    for term in terms:
         s = str(term)
         if res and not s.startswith("-"):
             res += "+"
@@ -85,3 +77,10 @@ def simplify(poly):
         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)