advent-of-code/2022/d21/ex1/ex1.py

81 lines
1.8 KiB
Python
Executable file

#!/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()