aoc/2022/11/solution.py

#!/usr/bin/python3
import sys
import re
import math
from dataclasses import dataclass
from typing import Callable

@dataclass
class Monkey:
    items: list
    operator: Callable[[int, int], int]
    has_num_arg: bool
    num_arg: int
    div_by: int
    to1: int
    to2: int
    reduce_worry_level: bool
    coeff: int = 0
    inspections: int = 0

    def turn(self):
        #print(self.items, '-> ', end='')
        throws = []
        for i in range(len(self.items)):
            # increase inspections
            self.inspections += 1

            # apply operation
            self.items[i] = self.operator(self.items[i], self.num_arg if self.has_num_arg else self.items[i])

            # keep worry levels small
            if self.reduce_worry_level:
                self.items[i] //= 3
            else:
                self.items[i] %= self.coeff

            # pick target
            target = self.to1 if (self.items[i] % self.div_by == 0) else self.to2

            # construct list transformations
            throws.append((self.items[i], target))
        self.items.clear()
        #print(throws)

        return throws

def add(a, b): return a + b
def mul(a, b): return a * b

def business(monkeys):
    for monkey in monkeys:
        throws = monkey.turn()
        for level, target in throws:
            monkeys[target].items.append(level)

def solution(rounds, reduce_worry_level):
    monkeys = []
    for line in sys.stdin:
        if 'Monkey' in line:
            starting_items = next(sys.stdin)
            operation_string = next(sys.stdin)
            test = next(sys.stdin)
            test_true = next(sys.stdin)
            test_false = next(sys.stdin)

            items = [int(i) for i in starting_items.strip().split(': ')[1].split(', ')]

            operation_params = sorted(re.match(r'^Operation: new = (old|-?[0-9]+) ([\*\+]) (old|-?[0-9]+)', operation_string.strip()).groups())
            operator = add if operation_params[0] == '+' else mul
            num_arg = int(operation_params[1]) if operation_params[1] != 'old' else 0
            has_num_arg = operation_params[1] != 'old'

            div_by = int(test.strip().split(' ')[-1])
            to1 = int(test_true.strip().split(' ')[-1])
            to2 = int(test_false.strip().split(' ')[-1])

            monkeys.append(Monkey(items, operator, has_num_arg, num_arg, div_by, to1, to2, reduce_worry_level))
    
    if not reduce_worry_level:
        coeff = 1
        for monkey in monkeys:
            coeff *= monkey.div_by
        for monkey in monkeys:
            monkey.coeff = coeff

    for _ in range(rounds):
        business(monkeys)

    inspections = []
    for monkey in monkeys:
        inspections.append(monkey.inspections)
    max1, max2 = sorted(inspections)[-2:]
    print(f'max. 2 inspection counts ({max1}, {max2}) yield total monkey business of {max1 * max2}')

if sys.argv[1] in '1':
    solution(20, True)
else:
    solution(10000, False)