advent-of-code/2022/day21.py
Mikaël Capelle 8a0412f926 Clean 2022.
2023-12-09 09:54:53 +01:00

108 lines
2.8 KiB
Python

import operator
import sys
from typing import Callable
def compute(monkeys: dict[str, int | tuple[str, str, str]], monkey: str) -> int:
value = monkeys[monkey]
if isinstance(value, int):
return value
else:
op: dict[str, Callable[[int, int], int]] = {
"+": operator.add,
"-": operator.sub,
"*": operator.mul,
"/": operator.floordiv,
}
value = op[value[1]](compute(monkeys, value[0]), compute(monkeys, value[2]))
monkeys[monkey] = value
return value
def invert(
monkeys: dict[str, int | tuple[str, str, str]], monkey: str, target: int
) -> dict[str, int | tuple[str, str, str]]:
"""
Revert the given mapping from monkey name to value or operation such that
the value from 'monkey' is computable by inverting operation until the root is
found.
Args:
monkeys: Dictionary of monkeys, that will be updated and returned.
monkey: Name of the monkey to start from.
target: Target value to set for the monkey that depends on root.
Returns:
The given dictionary of monkeys.
"""
monkeys = monkeys.copy()
depends: dict[str, str] = {}
for m, v in monkeys.items():
if isinstance(v, int):
continue
op1, _, op2 = v
assert op1 not in depends
assert op2 not in depends
depends[op1] = m
depends[op2] = m
invert_op = {"+": "-", "-": "+", "*": "/", "/": "*"}
current = monkey
while True:
dep = depends[current]
if dep == "root":
monkeys[current] = target
break
val = monkeys[dep]
assert not isinstance(val, int)
op1, ope, op2 = val
if op1 == current:
monkeys[current] = (dep, invert_op[ope], op2)
elif ope in ("+", "*"):
monkeys[current] = (dep, invert_op[ope], op1)
else:
monkeys[current] = (op1, ope, dep)
current = dep
return monkeys
lines = sys.stdin.read().splitlines()
monkeys: dict[str, int | tuple[str, str, str]] = {}
op_monkeys: set[str] = set()
for line in lines:
parts = line.split(":")
name = parts[0].strip()
try:
value = int(parts[1].strip())
monkeys[name] = value
except ValueError:
op1, ope, op2 = parts[1].strip().split()
monkeys[name] = (op1, ope, op2)
op_monkeys.add(name)
answer_1 = compute(monkeys.copy(), "root")
print(f"answer 1 is {answer_1}")
# assume the second operand of 'root' can be computed, and the first one depends on
# humn, which is the case is my input and the test input
p1, _, p2 = monkeys["root"] # type: ignore
answer_2 = compute(invert(monkeys, "humn", compute(monkeys.copy(), p2)), "humn")
print(f"answer 2 is {answer_2}")