day11.py

#!/usr/bin/env python3
from __future__ import annotations

import dataclasses
from textwrap import dedent
from typing import Callable

import aocd
import parse

test_data = """\
Monkey 0:
  Starting items: 79, 98
  Operation: new = old * 19
  Test: divisible by 23
    If true: throw to monkey 2
    If false: throw to monkey 3

Monkey 1:
  Starting items: 54, 65, 75, 74
  Operation: new = old + 6
  Test: divisible by 19
    If true: throw to monkey 2
    If false: throw to monkey 0

Monkey 2:
  Starting items: 79, 60, 97
  Operation: new = old * old
  Test: divisible by 13
    If true: throw to monkey 1
    If false: throw to monkey 3

Monkey 3:
  Starting items: 74
  Operation: new = old + 3
  Test: divisible by 17
    If true: throw to monkey 0
    If false: throw to monkey 1
""".split(
    "\n\n"
)

assert len(test_data) == 4
test_answer_one = 10605
test_answer_two = 2713310158

N_ROUNDS = 20


@dataclasses.dataclass(frozen=True)
class Monkey:
    items: list[int]
    operation: Callable[[int], int]
    test: Callable[[int], bool]
    p: int
    f: int


def parse_monkey(mstring) -> Monkey:
    result = parse.parse(
        """\
Monkey {id}:
  Starting items: {items}
  Operation: new = old {operation}
  Test: divisible by {divisor}
    If true: throw to monkey {t_m}
    If false: throw to monkey {f_m}""",
        mstring,
    )
    assert result, mstring

    def val(s) -> str:
        return result.named[s]

    return Monkey(
        items=[int(x) for x in val("items").split(", ")],
        operation=eval("lambda old: old " + val("operation")),
        test=lambda x: x % int(val("divisor")) == 0,
        p=int(val("t_m")),
        f=int(val("f_m")),
    )


def monkey_inspect(m: Monkey, divide: bool):
    pass_out = []
    fail_out = []

    for item in m.items:

        item = m.operation(item)
        if divide:
            item //= 3
        if m.test(item):
            pass_out.append(item)
        else:
            fail_out.append(item)
    m.items.clear()
    return pass_out, fail_out


def part_one(data: list[str]) -> int:

    monkey_business = []
    inspection_counter = []
    for mstr in data:
        m = parse_monkey(mstr)
        monkey_business.append(m)
        inspection_counter.append(0)

    # now do 20 rounds:
    for _ in range(N_ROUNDS):
        c = 0
        for monkey in monkey_business:
            inspection_counter[c] += len(monkey.items)
            p_out, f_out = monkey_inspect(monkey, True)
            monkey_business[monkey.p].items.extend(p_out)
            monkey_business[monkey.f].items.extend(f_out)
            c += 1

    # finc max two values
    a, b = sorted(inspection_counter, reverse=True)[:2]
    return a * b


def part_two(data: list[str]) -> int:
    monkey_business = []
    inspection_counter = []
    for mstr in data:
        m = parse_monkey(mstr)
        monkey_business.append(m)
        inspection_counter.append(0)

    # now do 10000 rounds:
    for _ in range(10000):
        c = 0
        for monkey in monkey_business:
            inspection_counter[c] += len(monkey.items)
            p_out, f_out = monkey_inspect(monkey, False)
            monkey_business[monkey.p].items.extend(p_out)
            monkey_business[monkey.f].items.extend(f_out)
            c += 1

    # finc max two values
    a, b = sorted(inspection_counter, reverse=True)[:2]
    return a * b


if __name__ == "__main__":

    part_one_ans = part_one(test_data)
    assert part_one_ans == test_answer_one, f"{part_one_ans=}, not {test_answer_one=}"

    real_data = aocd.get_data(day=11, year=2022).split("\n\n")
    print("part 1:", part_one(real_data))

    part_two_ans = part_two(test_data)
    assert part_two_ans == test_answer_two, f"{part_two_ans=}, not {test_answer_two=}"
    print("part 2:", part_two(real_data))