advent-of-code/2022/day11/rust/src/main.rs

#![warn(clippy::pedantic)]

use std::{
    cmp::Reverse,
    io::{stdin, Read},
    ops::Rem,
    str::FromStr,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Operand {
    Const(usize),
    Old,
}

impl Operand {
    fn evaluate(self, old: usize) -> usize {
        match self {
            Self::Old => old,
            Self::Const(c) => c,
        }
    }
}

impl FromStr for Operand {
    type Err = std::num::ParseIntError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "old" => Ok(Operand::Old),
            i => i.parse().map(Operand::Const),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Operation {
    Add(Operand),
    Multiply(Operand),
    Divide(Operand),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Test {
    Divisible(usize),
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct Item {
    value: usize,
}

impl Item {
    fn apply_operation(&mut self, op: Operation, modulo: Option<usize>) {
        match op {
            Operation::Add(x) => self.value += x.evaluate(self.value),
            Operation::Multiply(x) => self.value *= x.evaluate(self.value),
            Operation::Divide(x) => {
                assert!(modulo.is_none(), "can't divide under modulo arithmetic");
                self.value /= x.evaluate(self.value);
            }
        }

        if let Some(modulo) = modulo {
            self.value %= modulo;
        }
    }

    fn apply_test(&self, test: Test) -> bool {
        match test {
            Test::Divisible(n) => self.value.rem(n) == 0,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct Monkey {
    items: Vec<Item>,
    operation: Operation,
    test: Test,
    target_pass: usize,
    target_fail: usize,
    inspections: usize,
}

impl Monkey {
    fn throw_item(&mut self, divisor: Option<usize>, modulo: Option<usize>) -> (usize, Item) {
        self.inspections += 1;
        let mut item = self.items.remove(0);
        item.apply_operation(self.operation, modulo);
        if let Some(divisor) = divisor {
            item.apply_operation(Operation::Divide(Operand::Const(divisor)), modulo);
        }

        let target = if item.apply_test(self.test) {
            self.target_pass
        } else {
            self.target_fail
        };

        (target, item)
    }
}

impl FromStr for Monkey {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut lines = s.lines();
        lines.next().ok_or(())?; // Monkey n:

        let mut with_prefix = |prefix: &str| lines.next().ok_or(())?.strip_prefix(prefix).ok_or(());
        let items: Vec<_> = with_prefix("  Starting items: ")?
            .split(", ")
            .map(|i| {
                let value = i.parse().map_err(drop)?;
                Ok(Item { value })
            })
            .collect::<Result<_, _>>()?;

        let (operator, operand) = with_prefix("  Operation: new = old ")?
            .split_once(' ')
            .ok_or(())?;
        let operand = operand.parse().map_err(drop)?;
        let operation = match operator {
            "+" => Operation::Add(operand),
            "*" => Operation::Multiply(operand),
            _ => panic!("invalid operator {operator}"),
        };

        let test = Test::Divisible(
            with_prefix("  Test: divisible by ")?
                .parse()
                .map_err(drop)?,
        );

        let target_pass = with_prefix("    If true: throw to monkey ")?
            .parse()
            .map_err(drop)?;
        let target_fail = with_prefix("    If false: throw to monkey ")?
            .parse()
            .map_err(drop)?;

        Ok(Monkey {
            items,
            operation,
            test,
            target_pass,
            target_fail,
            inspections: 0,
        })
    }
}

fn play(monkeys: &mut [Monkey], rounds: usize, divisor: Option<usize>) -> usize {
    let modulo = match divisor {
        Some(_) => None,
        None => Some(
            monkeys
                .iter()
                .map(|m| {
                    let Test::Divisible(i) = m.test;
                    i
                })
                .product(),
        ),
    };

    for _round in 0..rounds {
        for i in 0..monkeys.len() {
            for _j in 0..monkeys[i].items.len() {
                let (target, item) = monkeys[i].throw_item(divisor, modulo);

                monkeys[target].items.push(item);
            }
        }
    }

    let mut levels: Vec<_> = monkeys.iter().map(|m| m.inspections).collect();
    levels.sort_unstable_by_key(|i| Reverse(*i));
    levels[0] * levels[1]
}

fn main() {
    let mut data = String::new();
    stdin().read_to_string(&mut data).unwrap();

    let mut monkeys: Vec<Monkey> = data
        .split("\n\n")
        .map(|monkey| monkey.parse().unwrap())
        .collect();

    println!("{:?}", play(&mut monkeys.clone(), 20, Some(3)));
    println!("{:?}", play(&mut monkeys, 10_000, None));
}