Day 18 (cleanup part 2)

day18/part2.py

@@ -1,48 +1,37 @@
 #!/usr/bin/env python3
-import re
 
 lines = (x.strip().replace(" ", "") for x in open("input.txt"))
+operators = {
+    "+": [2, lambda a, b: a+b],
+    "*": [1, lambda a, b: a*b],
+}
 
 
 def calc(a, b, pop):
-    return a+b if pop == "+" else a*b
+    return operators[pop][1](a, b)
 
 
 def pres(pop):
-    if pop == "+":
-        return 2
-    if pop == "*":
-        return 1
-    return 0
+    return operators.get(pop, [0])[0]
 
 
 def evaluate(line):
     stack = []
     postfix = [0]
-    pos = -1
     for char in "("+line+")":
-        pos += 1
-        if char == '(':
-            stack.append("(")
-        elif char == ')':
-            while True:
-                pop = stack.pop()
-                if pop == '(':
-                    break
-                a, b = postfix.pop(), postfix.pop()
-                postfix.append(calc(a, b, pop))
-        elif char in "0123456789":
+        if char.isdigit():
             postfix[-1] = postfix[-1] * 10 + int(char)
+        elif char == '(':
+            stack.append(char)
+        elif char == ')':
+            while pres(stack[-1]) > 0:
+                postfix.append(calc(postfix.pop(), postfix.pop(), stack.pop()))
+            stack.pop()  # remove (
         elif char in "*+":
-            while True:
-                pop = stack.pop()
-                if pres(pop) <= pres(char):
-                    stack.append(pop)
+            while pres(stack[-1]) > pres(char):
+                postfix.append(calc(postfix.pop(), postfix.pop(), stack.pop()))
             stack.append(char)
             postfix.append(0)
-                    break
-                a, b = postfix.pop(), postfix.pop()
-                postfix.append(calc(a, b, pop))
         else:
             raise SystemError("Invalid char: "+char)
     return postfix[0]