aoc2022/src/day11.rs

use std::collections::VecDeque;

use itertools::Itertools;
use nom::branch::alt;
use nom::bytes::complete::tag;
use nom::character::complete;
use nom::character::complete::{newline, space1};
use nom::multi::separated_list1;
use nom::sequence::{delimited, preceded};
use nom::{IResult, Parser};

pub fn process_part_1(input: &str) -> usize {
    let mut monkeys = parse(input).unwrap().1;

    // Sadly we are doing concurrent array indexing in rust with borrows
    for _round in 0..20 {
        for monkey_id in 0..monkeys.len() {
            while !monkeys[monkey_id].items.is_empty() {
                // I'm relying on rust's magic 'figure out when we can drop this to make it work'
                // here. This var can be dropped *just* before the last line of the block to make it
                // so we aren't borrowing mutably twice.
                let monkey = &mut monkeys[monkey_id];

                monkey.num_inspected += 1;
                let item = monkey.items.pop_front().unwrap();
                let item = monkey.operation.execute(item) / 3;
                let next_monkey_id = if item % monkey.test_mod == 0 {
                    monkey.target_true
                } else {
                    monkey.target_false
                };
                monkeys[next_monkey_id].items.push_back(item);
            }
        }
    }

    monkeys
        .iter()
        .map(|monkey| monkey.num_inspected)
        .sorted()
        .rev()
        .take(2)
        .product()
}

pub fn process_part_2(input: &str) -> usize {
    let mut monkeys = parse(input).unwrap().1;

    // Since we only care for divisibility in tests, we can just modulo every item with
    // the LCM of all tests. This allows all tests to check for eligibility without actually
    // letting the number grow huge.
    let lcm = monkeys
        .iter()
        .map(|x| x.test_mod)
        .reduce(num::integer::lcm)
        .unwrap();

    for _round in 0..10000 {
        for monkey_id in 0..monkeys.len() {
            while !monkeys[monkey_id].items.is_empty() {
                let monkey = &mut monkeys[monkey_id];
                monkey.num_inspected += 1;

                let item = monkey.items.pop_front().unwrap();
                let item = monkey.operation.execute(item) % lcm;
                let next_monkey_id = if item % monkey.test_mod == 0 {
                    monkey.target_true
                } else {
                    monkey.target_false
                };
                monkeys[next_monkey_id].items.push_back(item);
            }
        }
    }

    monkeys
        .iter()
        .map(|monkey| monkey.num_inspected)
        .sorted()
        .rev()
        .take(2)
        .product()
}

#[derive(Debug)]
enum Operation {
    Add(u64),
    Multiply(u64),
    MultiplySelf,
}

impl Operation {
    fn execute(&self, item: u64) -> u64 {
        match self {
            Operation::Add(x) => item + x,
            Operation::Multiply(x) => item * x,
            Operation::MultiplySelf => item * item,
        }
    }
}

#[derive(Debug)]
struct Monkey {
    items: VecDeque<u64>,
    operation: Operation,
    test_mod: u64,
    target_true: usize,
    target_false: usize,
    num_inspected: usize,
}

fn parse(input: &str) -> IResult<&str, Vec<Monkey>> {
    separated_list1(newline, parse_monkey)(input)
}

fn parse_monkey(input: &str) -> IResult<&str, Monkey> {
    let (input, _id) = delimited(tag("Monkey "), complete::u64, tag(":\n"))(input)?;
    let (input, items) = preceded(
        space1,
        delimited(
            tag("Starting items: "),
            separated_list1(tag(", "), complete::u64),
            newline,
        ),
    )(input)?;
    let (input, operation) = delimited(space1, parse_operation, newline)(input)?;
    let (input, test_mod) = preceded(
        space1,
        delimited(tag("Test: divisible by "), complete::u64, newline),
    )(input)?;
    let (input, target_true) = preceded(
        space1,
        delimited(tag("If true: throw to monkey "), complete::u64, newline),
    )(input)?;
    let (input, target_false) = preceded(
        space1,
        delimited(tag("If false: throw to monkey "), complete::u64, newline),
    )(input)?;

    Ok((
        input,
        Monkey {
            items: VecDeque::from(items),
            operation,
            test_mod,
            target_true: target_true as usize,
            target_false: target_false as usize,
            num_inspected: 0,
        },
    ))
}

fn parse_operation(input: &str) -> IResult<&str, Operation> {
    let (input, _) = tag("Operation: new = old ")(input)?;
    alt((
        tag("* old").map(|_| Operation::MultiplySelf),
        preceded(tag("+ "), complete::u64).map(Operation::Add),
        preceded(tag("* "), complete::u64).map(Operation::Multiply),
    ))(input)
}

#[cfg(test)]
mod tests {
    use super::*;

    const INPUT: &str = "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
";

    #[test]
    fn day1() {
        assert_eq!(process_part_1(INPUT), 10605);
    }

    #[test]
    fn day2() {
        assert_eq!(process_part_2(INPUT), 2713310158);
    }
}