advent-of-code/2022/d11/ex2/ex2.py

96 lines
2.6 KiB
Python
Executable file

#!/usr/bin/env python
import dataclasses
import math
import sys
from collections import Counter, deque
from typing import Literal, Optional
@dataclasses.dataclass
class Operation:
lhs: Optional[int]
operator: Literal["+", "*"]
rhs: Optional[int]
def __call__(self, old: int) -> int:
lhs = old if self.lhs is None else self.lhs
rhs = old if self.rhs is None else self.rhs
if self.operator == "*":
return lhs * rhs
if self.operator == "+":
return lhs + rhs
assert False
@classmethod
def from_input(cls, input: str) -> "Operation":
assert input.startswith(" Operation: new = ")
lhs, op, rhs = input.split()[-3:]
assert op in ("+", "*") # Sanity check
return cls(
None if lhs == "old" else int(lhs),
op, # type: ignore
None if rhs == "old" else int(rhs),
)
@dataclasses.dataclass
class Monkey:
items: deque[int]
operation: Operation
test_divisor: int
transfer: dict[bool, int]
@classmethod
def from_input(cls, input: list[str]) -> "Monkey":
# Sanity checks
assert input[0].startswith("Monkey ")
assert "divisible by" in input[3]
assert "true" in input[4]
assert "false" in input[5]
items = deque(
int(n) for n in input[1].removeprefix(" Starting items: ").split(",")
)
operation = Operation.from_input(input[2])
divisor = int(input[3].split()[-1])
transfer = {
True: int(input[4].split()[-1]),
False: int(input[5].split()[-1]),
}
return Monkey(items, operation, divisor, transfer)
def solve(input: list[str]) -> int:
def do_round(monkeys: list[Monkey], counts: dict[int, int]) -> None:
modulo = math.prod(monkey.test_divisor for monkey in monkeys)
for i, monkey in enumerate(monkeys):
counts[i] += len(monkey.items)
while monkey.items:
item = monkey.items.popleft()
item = monkey.operation(item)
item %= modulo # Keep a reasonable range on this value
target = monkey.transfer[(item % monkey.test_divisor) == 0]
monkeys[target].items.append(item)
monkeys = [Monkey.from_input(monkey_spec.splitlines()) for monkey_spec in input]
counts: Counter[int] = Counter()
for _ in range(10_000):
do_round(monkeys, counts)
((_, a), (_, b)) = counts.most_common(2)
return a * b
def main() -> None:
input = sys.stdin.read().split("\n\n")
print(solve(input))
if __name__ == "__main__":
main()