21 entirely in python (with z3)

1 files changed, 8 insertions, 25 deletions

diff --git a/21.py b/21.py
index 302438d..1a44a55 100644
--- a/21.py
+++ b/21.py
@@ -1,7 +1,8 @@
 from functools import cache
-from operator import add, sub, mul, ifloordiv
+from operator import add, sub, mul, truediv, eq
 from subprocess import run, PIPE
 from utils import open_day
+from z3 import Int, ArithRef, solve
 
 monkeys = {}
 with open_day(21) as f:
@@ -12,33 +13,15 @@ with open_day(21) as f:
             expr = int(expr[0])
         monkeys[target] = expr
 
-ops = { '+': add, '-': sub, '*': mul, '/': ifloordiv }
+ops = { '+': add, '-': sub, '*': mul, '/': truediv, '==': eq }
 
-@cache
 def eval_monkey(m):
     expr = monkeys[m]
-    if isinstance(expr, int):
+    if isinstance(expr, int) or isinstance(expr, ArithRef):
         return expr
     return ops[expr[1]](eval_monkey(expr[0]), eval_monkey(expr[2]))
 
-print(eval_monkey('root'))
-
-@cache
-def to_sexp(m):
-    if m == 'humn': return m
-    expr = monkeys[m]
-    if isinstance(expr, int):
-        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'
-with open(filename, 'w') as f:
-    print(f'(declare-const humn Int)', file=f)
-    root = monkeys['root']
-    print(f'(assert (= {to_sexp(root[0])} {to_sexp(root[2])}))', file=f)
-    print('(check-sat)\n(get-value (humn))', file=f)
-r = run(['z3', filename], stdout=PIPE)
-print(r.stdout.splitlines()[1].decode().strip('()').split()[1])
+print(int(eval_monkey('root')))
+monkeys['humn'] = Int('humn')
+monkeys['root'][1] = '=='
+solve(eval_monkey('root'))