Day 19: Not Enough Minerals

Iterated Width

Width-based Planning is a search method that was, among other things, used successfully to play ATARI games in almost-realtime. The idea is to use a Breadth-First Search with a novelty pruning rule.
Each state is associated with features or atoms that are true for that state (e.g. "geodes=7"). The algorithm is called with increasing widths starting from 1. New states are pruned if there is no n-tuple of features (of size width) that has not already been seen during the search.

So if a state has the atoms 'A', 'B' and 'C', then IW(1) prunes the state if A, B and C have been seen before, and IW(2) prunes a state if (A, B), (A, C), and (B, C) have all been seen before. It is known that many problems have a low intrinsic width, so they are solvable using few calls to the IW search procedure, even if they have many different features.

I elected to use the complete search state of the problem (i.e. one atom for each robot and ore count, as well as the time remaining, and the current robot in production).

Apart from not knowing whether it would work at all, the main problem was that IW is an algorithm that is used for planning when you know the goal state, and here we didn't, so there was no way to know when to stop increasing the width parameter. This is probably highly unsound, but I stopped incrementing once the same result was returned twice in succession, and that worked.

The running time for both parts together is between 15 and 20 seconds, and that's with using multiple cores, so that could be better...

use std::collections::{HashSet, VecDeque};
use std::hash::{Hash, Hasher};

use itertools::Itertools;
use rayon::prelude::*;
use wyhash::WyHash; // For hashing. Faster than rust's default SIPhash

use parser::parse;

mod parser; // Parses the puzzle input (omitted)

#[derive(Clone, Debug)]
pub struct Blueprint {
    number: usize,
    ore_robot_ore_cost: usize,
    clay_robot_ore_cost: usize,
    obsidian_robot_ore_cost: usize,
    obsidian_robot_clay_cost: usize,
    geode_robot_ore_cost: usize,
    geode_robot_obsidian_cost: usize,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
enum Atom {
    Time(usize),
    Ore(usize),
    Clay(usize),
    Obsidian(usize),
    Geodes(usize),
    OreRobot(usize),
    ClayRobot(usize),
    ObsidianRobot(usize),
    GeodeRobot(usize),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SearchState {
    time_remaining: usize,
    ore: usize,
    clay: usize,
    obsidian: usize,
    geodes: usize,
    ore_robots: usize,
    clay_robots: usize,
    obsidian_robots: usize,
    geode_robots: usize,
    currently_building: Option<Atom>, // This is a bit unclean, but eh.
}

impl SearchState {
    // Advance the simulation by one minute.
    // This is quite inelegant, but I'm out of time and won't refactor it now.
    pub fn tick(&self) -> Self {
        let time_remaining = self.time_remaining - 1;

        let ore = self.ore + self.ore_robots;
        let clay = self.clay + self.clay_robots;
        let obsidian = self.obsidian + self.obsidian_robots;
        let geodes = self.geodes + self.geode_robots;

        match self.currently_building {
            Some(Atom::OreRobot(_)) => Self {
                time_remaining,
                ore,
                clay,
                obsidian,
                geodes,
                ore_robots: self.ore_robots + 1,
                currently_building: None,
                ..*self
            },

            Some(Atom::ClayRobot(_)) => Self {
                time_remaining,
                ore,
                clay,
                obsidian,
                geodes,
                clay_robots: self.clay_robots + 1,
                currently_building: None,
                ..*self
            },

            Some(Atom::ObsidianRobot(_)) => Self {
                time_remaining,
                ore,
                clay,
                obsidian,
                geodes,
                obsidian_robots: self.obsidian_robots + 1,
                currently_building: None,
                ..*self
            },

            Some(Atom::GeodeRobot(_)) => Self {
                time_remaining,
                ore,
                clay,
                obsidian,
                geodes,
                geode_robots: self.geode_robots + 1,
                currently_building: None,
                ..*self
            },

            None => Self {
                time_remaining,
                ore,
                clay,
                obsidian,
                geodes,
                ..*self
            },

            _ => unreachable!(),
        }
    }

    // The atoms of a search state
    pub fn atoms(&self) -> impl Iterator<Item = Atom> + '_ {
        let mut atoms = vec![
            Atom::Time(self.time_remaining),
            Atom::Ore(self.ore),
            Atom::Clay(self.clay),
            Atom::Obsidian(self.obsidian),
            Atom::Geodes(self.geodes),
            Atom::OreRobot(self.ore_robots),
            Atom::ClayRobot(self.clay_robots),
            Atom::ObsidianRobot(self.obsidian_robots),
            Atom::GeodeRobot(self.geode_robots),
        ];

        if self.currently_building.is_some() {
            atoms.push(self.currently_building.unwrap());
        }

        atoms.into_iter()
    }

    // Buy the robots.
    //
    // Each function checks if the purchase is valid, and if so, returns a new
    // search state with the purchase made.
    //
    // There are minor pruning rules here. For example, we don't buy a
    // obsidian bot if we already have enough bots so we can afford to buy
    // a geode bot every turn (wrt obsidian cost).

    pub fn buy_geode_bot(&self, blueprint: &Blueprint) -> Option<Self> {
        if self.time_remaining > 1
            && self.obsidian >= blueprint.geode_robot_obsidian_cost
            && self.ore >= blueprint.geode_robot_ore_cost
        {
            Some(Self {
                geode_robots: self.geode_robots,
                ore: self.ore - blueprint.geode_robot_ore_cost,
                obsidian: self.obsidian - blueprint.geode_robot_obsidian_cost,
                currently_building: Some(Atom::GeodeRobot(1)),
                ..*self
            })
        } else {
            None
        }
    }

    pub fn buy_obsidian_bot(&self, blueprint: &Blueprint) -> Option<Self> {
        if self.time_remaining > 1
            && self.clay >= blueprint.obsidian_robot_clay_cost
            && self.ore >= blueprint.obsidian_robot_ore_cost
            && self.obsidian_robots < blueprint.geode_robot_obsidian_cost
        {
            Some(Self {
                obsidian_robots: self.obsidian_robots,
                ore: self.ore - blueprint.obsidian_robot_ore_cost,
                clay: self.clay - blueprint.obsidian_robot_clay_cost,
                currently_building: Some(Atom::ObsidianRobot(1)),
                ..*self
            })
        } else {
            None
        }
    }

    pub fn buy_clay_bot(&self, blueprint: &Blueprint) -> Option<Self> {
        if self.time_remaining > 1
            && self.ore >= blueprint.clay_robot_ore_cost
            && self.clay_robots < blueprint.obsidian_robot_clay_cost
        {
            Some(Self {
                clay_robots: self.clay_robots,
                ore: self.ore - blueprint.clay_robot_ore_cost,
                currently_building: Some(Atom::ClayRobot(1)),
                ..*self
            })
        } else {
            None
        }
    }

    pub fn buy_ore_bot(&self, blueprint: &Blueprint) -> Option<Self> {
        let max_cost = blueprint
            .ore_robot_ore_cost
            .max(blueprint.clay_robot_ore_cost)
            .max(blueprint.obsidian_robot_ore_cost)
            .max(blueprint.geode_robot_ore_cost);

        if self.time_remaining > blueprint.ore_robot_ore_cost + 1
            && self.ore >= blueprint.ore_robot_ore_cost
            && self.ore_robots < max_cost
        {
            Some(Self {
                ore_robots: self.ore_robots,
                ore: self.ore - blueprint.ore_robot_ore_cost,
                currently_building: Some(Atom::OreRobot(1)),
                ..*self
            })
        } else {
            None
        }
    }

    // Prune checks if the state should be pruned according to the IW(k) novelty
    // rule. If so, it returns None, otherwise it returns the state itself and a
    // Vec of the atoms that were newly seen.
    //
    // Again, this could probably done more elegantly.
    //
    pub fn prune(self, true_atoms: &HashSet<u64>, width: usize) -> Option<(Self, Vec<u64>)> {
        let mut prune = true;
        let mut new_atoms = Vec::new();

        for atom in self.atoms().combinations(width) {
            let mut h = WyHash::default();
            atom.hash(&mut h);
            let atom_hash = h.finish();

            if true_atoms.contains(&atom_hash) {
                continue;
            } else {
                prune = false;
                new_atoms.push(atom_hash);
            }
        }

        if prune {
            return None;
        }

        Some((self, new_atoms))
    }
}

// Given a blueprint, calculate how many geodes we can open in the given time
fn calculate_openable_geodes(blueprint: Blueprint, time_remaining: usize, width: usize) -> usize {
    let mut q = VecDeque::new();
    let mut true_atoms: HashSet<u64> = HashSet::new();
    let mut best = 0;

    let initial_state = SearchState {
        time_remaining,
        ore: 0,
        clay: 0,
        obsidian: 0,
        geodes: 0,
        ore_robots: 1,
        clay_robots: 0,
        obsidian_robots: 0,
        geode_robots: 0,
        currently_building: None,
    };

    q.push_back(initial_state);

    while let Some(state) = q.pop_front() {
        if state.time_remaining == 0 {
            best = best.max(state.geodes);
            continue;
        }

        let next_state = state.tick();

        // Always try to buy a geode robot if possible, not doing so does not make sense
        if let Some(buy_geode_bot) = next_state.buy_geode_bot(&blueprint) {
            if let Some((novel, new_atoms)) = buy_geode_bot.prune(&true_atoms, width) {
                q.push_back(novel);
                true_atoms.extend(new_atoms);
            }
        } else {
            // If we can't buy a geode cracker, try the other options
            if let Some(buy_obsidian_bot) = next_state.buy_obsidian_bot(&blueprint) {
                if let Some((novel, new_atoms)) = buy_obsidian_bot.prune(&true_atoms, width) {
                    q.push_back(novel);
                    true_atoms.extend(new_atoms);
                }
            }

            if let Some(buy_clay_bot) = next_state.buy_clay_bot(&blueprint) {
                if let Some((novel, new_atoms)) = buy_clay_bot.prune(&true_atoms, width) {
                    q.push_back(novel);
                    true_atoms.extend(new_atoms);
                }
            }

            if let Some(buy_ore_bot) = next_state.buy_ore_bot(&blueprint) {
                if let Some((novel, new_atoms)) = buy_ore_bot.prune(&true_atoms, width) {
                    q.push_back(novel);
                    true_atoms.extend(new_atoms);
                }
            }

            // Do nothing
            if let Some((novel, new_atoms)) = next_state.prune(&true_atoms, width) {
                q.push_back(novel);
                true_atoms.extend(new_atoms);
            }
        }
    }

    best
}

// Repeatedly increase the IW width until we get the same result twice
// or the width exceeds 4.
fn evaluate_blueprint(blueprint: Blueprint, time: usize) -> usize {
    let mut width = 1;
    let mut prev_result = None;

    while width < 4 {
        let result = calculate_openable_geodes(blueprint.clone(), time, width);

        if Some(result) == prev_result {
            break;
        }

        prev_result = Some(result);
        width += 1;
    }

    prev_result.unwrap()
}

fn part1(blueprints: Vec<Blueprint>) -> usize {
    blueprints
        .into_par_iter()
        .map(|bp| bp.number * evaluate_blueprint(bp, 24))
        .sum::<usize>()
}

fn part2(blueprints: Vec<Blueprint>) -> usize {
    blueprints
        .into_par_iter()
        .map(|bp| evaluate_blueprint(bp, 32))
        .reduce(|| 1, |a, b| a * b)
}

fn input_text() -> &'static str {
    include_str!("../input.txt")
}

fn main() {
    dbg!(part1(parse(input_text())));

    let p2_blueprints = parse(input_text());

    dbg!(part2(
        p2_blueprints.iter().take(3).cloned().collect::<Vec<_>>()
    ));
}