advent/y2022/p11.py

import operator
import re
from collections import namedtuple
from math import prod


def make_monkey(description):
    relevant = re.findall(r'\d+|old|\*|\+', description)
    parsed = [int(n) if n.isdigit() else n for n in relevant]
    _, *items, _, op, y, div, aa, bb = parsed

    op = {'*': operator.mul, '+': operator.add}[op]
    grow_item = lambda x: op(x, x if y == 'old' else y)
    find_target = lambda x: aa if x % div == 0 else bb
    return Monkey(items, grow_item, find_target)


def play(monkeys, n_rounds, shrink):
    counts = [0 for _ in monkeys]
    for _ in range(n_rounds):
        for idx, monkey in enumerate(monkeys):
            while monkey.items:
                counts[idx] += 1
                item = monkey.items.pop()
                item = monkey.grow_item(item)
                item = shrink(item)
                idx = monkey.find_target(item)
                monkeys[idx].items.append(item)
    return prod(sorted(counts)[-2:])


text = open(0).read()
Monkey = namedtuple('Monkey', ['items', 'grow_item', 'find_target'])

monkeys = [make_monkey(descr) for descr in text.strip().split('\n\n')]
ans1 = play(monkeys, 20, lambda x: x // 3)
print(ans1)

lim = prod(map(int, re.findall(r'divisible by (\d+)', text)))
monkeys = [make_monkey(descr) for descr in text.strip().split('\n\n')]
ans2 = play(monkeys, 10_000, lambda x: x % lim)
print(ans2)