146 lines
4.3 KiB
Python
146 lines
4.3 KiB
Python
from __future__ import annotations
|
|
from typing import *
|
|
from aocpy import BaseChallenge
|
|
from enum import Enum
|
|
import re
|
|
|
|
|
|
class Operator(Enum):
|
|
ADD = "+"
|
|
SUB = "-"
|
|
DIV = "/"
|
|
MULT = "*"
|
|
|
|
def inverse(self) -> Operator:
|
|
if self.value == self.ADD.value:
|
|
return self.SUB
|
|
elif self.value == self.SUB.value:
|
|
return self.ADD
|
|
elif self.value == self.MULT.value:
|
|
return self.DIV
|
|
else:
|
|
return self.MULT
|
|
|
|
def calc(self, a_val, b_val) -> float:
|
|
res = 0
|
|
if self.value == self.ADD.value:
|
|
res = a_val + b_val
|
|
elif self.value == self.SUB.value:
|
|
res = a_val - b_val
|
|
elif self.value == self.MULT.value:
|
|
res = a_val * b_val
|
|
elif self.value == self.DIV.value:
|
|
res = a_val / b_val
|
|
|
|
return res
|
|
|
|
|
|
Monkey = Union[int, Tuple[str, Operator, str]]
|
|
Monkies = Dict[str, Monkey]
|
|
|
|
|
|
number_re = re.compile(r"([a-z]{4}): (\d+)")
|
|
operation_re = re.compile(r"([a-z]{4}): ([a-z]{4}) ([+\-/*]) ([a-z]{4})")
|
|
|
|
|
|
def parse(instr: str) -> Monkies:
|
|
res: Dict[str, Monkey] = {}
|
|
for line in instr.strip().splitlines():
|
|
if x := number_re.match(line):
|
|
name, num_str = x.groups()
|
|
num = int(num_str)
|
|
res[name] = num
|
|
elif x := operation_re.match(line):
|
|
name, a, op_str, b = x.groups()
|
|
res[name] = (a, Operator(op_str), b)
|
|
return res
|
|
|
|
|
|
def evaluate_monkey(monkey: Monkey, monkies: Monkies) -> float:
|
|
if type(monkey) == int:
|
|
return monkey
|
|
|
|
elif type(monkey) == tuple:
|
|
(a_name, op, b_name) = monkey
|
|
a_val = evaluate_monkey(monkies[a_name], monkies)
|
|
b_val = evaluate_monkey(monkies[b_name], monkies)
|
|
return op.calc(a_val, b_val)
|
|
|
|
raise ValueError(f"impossible type {type(monkey)}")
|
|
|
|
|
|
def does_branch_contain(target: str, monkey: Monkey, monkies: Monkies) -> bool:
|
|
if type(monkey) == int:
|
|
return False
|
|
|
|
elif type(monkey) == tuple:
|
|
(a_name, _, b_name) = monkey
|
|
if a_name == target or b_name == target:
|
|
return True
|
|
|
|
if does_branch_contain(target, monkies[a_name], monkies):
|
|
return True
|
|
return does_branch_contain(target, monkies[b_name], monkies)
|
|
|
|
raise ValueError(f"impossible type {type(monkey)}")
|
|
|
|
|
|
class Challenge(BaseChallenge):
|
|
@staticmethod
|
|
def one(instr: str) -> int:
|
|
inp = parse(instr)
|
|
return int(evaluate_monkey(inp["root"], inp))
|
|
|
|
@staticmethod
|
|
def two(instr: str) -> int:
|
|
inp = parse(instr)
|
|
|
|
root = inp["root"]
|
|
assert type(root) == tuple
|
|
|
|
(a_name, _, b_name) = root
|
|
|
|
if does_branch_contain("humn", inp[a_name], inp):
|
|
branch_with_human = inp[a_name]
|
|
target_value = evaluate_monkey(inp[b_name], inp)
|
|
elif does_branch_contain("humn", inp[b_name], inp):
|
|
branch_with_human = inp[b_name]
|
|
target_value = evaluate_monkey(inp[a_name], inp)
|
|
else:
|
|
raise ValueError("neither root branch contains humn")
|
|
|
|
acc = target_value
|
|
cursor = branch_with_human
|
|
while True:
|
|
(a_name, op, b_name) = cursor
|
|
|
|
if a_name == "humn" or does_branch_contain("humn", inp[a_name], inp):
|
|
x = evaluate_monkey(inp[b_name], inp)
|
|
is_first_arg = False
|
|
cursor = inp[a_name]
|
|
elif b_name == "humn" or does_branch_contain("humn", inp[b_name], inp):
|
|
x = evaluate_monkey(inp[a_name], inp)
|
|
is_first_arg = True
|
|
cursor = inp[b_name]
|
|
else:
|
|
raise ValueError("neither branch contains humn")
|
|
|
|
if op == Operator.ADD or op == Operator.MULT:
|
|
acc = op.inverse().calc(acc, x)
|
|
elif op == Operator.SUB:
|
|
if is_first_arg:
|
|
acc = op.calc(-acc, -x)
|
|
else:
|
|
acc = op.inverse().calc(acc, x)
|
|
elif op == Operator.DIV:
|
|
if is_first_arg:
|
|
acc = op.calc(1 / acc, 1 / x)
|
|
else:
|
|
acc = op.inverse().calc(acc, x)
|
|
else:
|
|
raise ValueError(f"invalid operator {op}")
|
|
|
|
if a_name == "humn" or b_name == "humn":
|
|
break
|
|
|
|
return int(acc)
|