diff options
-rw-r--r-- | 21.py | 33 |
1 files changed, 8 insertions, 25 deletions
@@ -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')) |