21 simplify z3 program

1 files changed, 11 insertions, 10 deletions

diff --git a/21.py b/21.py
index d79cf0a..aff0b84 100644
--- a/21.py
+++ b/21.py
@@ -27,18 +27,19 @@ def eval_monkey(m):
 
 print(eval_monkey('root'))
 
+@cache
+def to_sexp(m):
+    if m == 'humn': return m
+    expr = monkeys[m]
+    if isinstance(expr, int):
+        return str(expr)
+    return f'({expr[1]} {to_sexp(expr[0])} {to_sexp(expr[2])})'
+
 filename = '21.z3'
 with open(filename, 'w') as f:
-    f.write(''.join(f'(declare-const {monkey} Int)\n' for monkey in monkeys.keys()))
-    for k, v in monkeys.items():
-        if k == 'humn': continue
-        if k == 'root':
-            print(f'(assert (= {v[0]} {v[2]}))', file=f)
-            continue
-        if isinstance(v, int):
-            print(f'(assert (= {k} {v}))', file=f)
-        else:
-            print(f'(assert (= {k} ({v[1]} {v[0]} {v[2]})))', file=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])