21

1 files changed, 7 insertions, 8 deletions

diff --git a/21.py b/21.py
index aff0b84..302438d 100644
--- a/21.py
+++ b/21.py
@@ -1,6 +1,7 @@
 from functools import cache
-from utils import open_day
+from operator import add, sub, mul, ifloordiv
 from subprocess import run, PIPE
+from utils import open_day
 
 monkeys = {}
 with open_day(21) as f:
@@ -11,12 +12,7 @@ with open_day(21) as f:
             expr = int(expr[0])
         monkeys[target] = expr
 
-ops = {
-    '+': lambda a, b: a + b,
-    '-': lambda a, b: a - b,
-    '*': lambda a, b: a * b,
-    '/': lambda a, b: a // b,
-}
+ops = { '+': add, '-': sub, '*': mul, '/': ifloordiv }
 
 @cache
 def eval_monkey(m):
@@ -32,7 +28,10 @@ def to_sexp(m):
     if m == 'humn': return m
     expr = monkeys[m]
     if isinstance(expr, int):
-        return str(expr)
+        return expr
+    a, b = to_sexp(expr[0]), to_sexp(expr[2])
+    if isinstance(a, int) and isinstance(b, int):
+        return ops[expr[1]](a, b)
     return f'({expr[1]} {to_sexp(expr[0])} {to_sexp(expr[2])})'
 
 filename = '21.z3'