diff --git a/2022/d21/ex1/ex1.py b/2022/d21/ex1/ex1.py new file mode 100755 index 0000000..7c12836 --- /dev/null +++ b/2022/d21/ex1/ex1.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python + +import dataclasses +import enum +import operator +import sys +from typing import Optional + + +class Operator(str, enum.Enum): + ADD = "+" + SUB = "-" + MUL = "*" + DIV = "/" + + def __call__(self, lhs: int, rhs: int) -> int: + OPERATIONS = { + self.ADD: operator.add, + self.SUB: operator.sub, + self.MUL: operator.mul, + self.DIV: operator.floordiv, + } + return OPERATIONS[self](lhs, rhs) + + +class Monkey: + def get_value(self, monkeys: dict[str, "Monkey"]) -> int: + raise NotImplemented + + +@dataclasses.dataclass +class YellerMonkey(Monkey): + value: int + + def get_value(self, monkeys: dict[str, "Monkey"]) -> int: + return self.value + + +@dataclasses.dataclass +class MathMonkey(Monkey): + lhs: str + op: Operator + rhs: str + _value: Optional[int] = dataclasses.field(default=None, init=False) + + def get_value(self, monkeys: dict[str, "Monkey"]) -> int: + if self._value is None: + self._value = self.op( + monkeys[self.lhs].get_value(monkeys), + monkeys[self.rhs].get_value(monkeys), + ) + return self._value + + +def solve(input: list[str]) -> int: + def parse_monkey(input: str) -> tuple[str, Monkey]: + name, value = input.split(": ") + + monkey: Monkey + match value.split(): + case [lhs, op, rhs]: + monkey = MathMonkey(lhs, Operator(op), rhs) + case [n]: + monkey = YellerMonkey(int(n)) + case _: + assert False # Sanity check + + return name, monkey + + monkeys = dict(map(parse_monkey, input)) + return monkeys["root"].get_value(monkeys) + + +def main() -> None: + input = sys.stdin.read().splitlines() + print(solve(input)) + + +if __name__ == "__main__": + main()