96 lines
2.6 KiB
Python
96 lines
2.6 KiB
Python
|
#!/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()
|