adventofcode2022/day11/sol.py

from collections import namedtuple
from copy import deepcopy
import math

Monkey = namedtuple('Monkey', 'index, op, divisor, target')
def parse(input):
    monkeys = []
    items = {}
    for i, chunk in enumerate(input.split("\n\n")):
        lines = chunk.strip().split("\n")
        assert 'Starting items:' in lines[1]
        assert 'Operation: new =' in lines[2]
        assert 'Test: divisible by' in lines[3]
        assert 'If true: throw to' in lines[4]
        assert 'If false: throw to' in lines[5]

        items[i] = [int(x.strip()) for x in lines[1].split(":")[1].split(",")]
        op = eval("lambda old: " + lines[2].split("=")[1])
        divisor = int(lines[3].split()[-1])
        target = [
            int(lines[5].split()[-1]),
            int(lines[4].split()[-1]),
        ]
        monkeys.append(Monkey(i, op, divisor, target))
    return monkeys, items

def play(monkeys, items, rounds=1, N=None):
    throws = {m.index: 0 for m in monkeys}
    for _ in range(rounds):
        for m in monkeys:
            for x in items[m.index]:
                if N:
                    nx = m.op(x) % N
                else:
                    nx = m.op(x) // 3
                trg = m.target[nx % m.divisor == 0]
                #print(f'{i}: {x} -> {nx}. throwing to {trg}')
                items[trg].append(nx)
            throws[m.index] += len(items[m.index])
            items[m.index] = []
    show(items)
    return throws

def show(items):
    for k in items:
        print(k, items[k])

def monkeybusiness(throws):
    a, b = sorted(throws.values())[-2:]
    return a*b

def lcm(ints):
    n = 1
    for x in ints:
        n *= x // math.gcd(n, x)
    return n

def main():
    monkeys, items = parse(open("input").read())
    #print(monkeys)
    show(items)

    throws = play(monkeys, deepcopy(items), rounds=20)
    print(monkeybusiness(throws))

    N = lcm(m.divisor for m in monkeys)
    throws = play(monkeys, deepcopy(items), rounds=10000, N=N)
    print(monkeybusiness(throws))

main()