Advent of Code 2021 - Day 23: Amphipod

amphipod.rs

fn main() {
    let burrow = Burrow::<8>::from(PART1);
    let result = burrow.energy_to_organise().unwrap();
    println!("Energy required to organise partial map: {:?}", result);

    let burrow = Burrow::<16>::from(PART2);
    let result = burrow.energy_to_organise().unwrap();
    println!("Energy required to organise full map: {:?}", result);
}

#[derive(Copy, Clone, Debug, PartialEq)]
enum Location {
    Wall,
    Corridor,
    Threshold,
    Room(Species),
}

#[derive(Copy, Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
enum Species {
    Amber,
    Bronze,
    Copper,
    Desert,
}

impl Species {
    fn energy_consumption(&self) -> u64 {
        match self {
            Self::Amber => 1,
            Self::Bronze => 10,
            Self::Copper => 100,
            Self::Desert => 1000,
        }
    }
}

#[derive(Copy, Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
struct Amphipod {
    species: Species,
    x: usize,
    y: usize,
}

impl Amphipod {
    fn possible_destinations<const L: usize>(
        &self,
        amphipods: [Amphipod; L],
        map: &[Vec<Location>],
    ) -> Vec<(usize, usize, u64)> {
        let mut destinations = Vec::new();
        let mut steps_made_leaving_the_room = 0;

        // Move out of the room if needed
        if self.y != 1 {
            let mut y = self.y - 1;

            while y > 0 && !amphipods.iter().any(|a| a.x == self.x && a.y == y) {
                steps_made_leaving_the_room += 1;
                y -= 1;
            }

            if y != 0 {
                return destinations;
            }
        }

        destinations.append(&mut self.possible_destinations_from_corridor(
            -1,
            steps_made_leaving_the_room,
            amphipods,
            map,
        ));
        destinations.append(&mut self.possible_destinations_from_corridor(
            1,
            steps_made_leaving_the_room,
            amphipods,
            map,
        ));

        destinations
    }

    fn possible_destinations_from_corridor<const L: usize>(
        &self,
        direction: i64,
        steps_made_leaving_the_room: u64,
        amphipods: [Amphipod; L],
        map: &[Vec<Location>],
    ) -> Vec<(usize, usize, u64)> {
        let mut destinations = Vec::new();

        let mut steps_made = steps_made_leaving_the_room;
        let mut x = (self.x as i64 + direction) as usize;

        while map[1][x] != Location::Wall && !amphipods.iter().any(|a| a.x == x && a.y == 1) {
            steps_made += 1;

            if map[1][x] != Location::Threshold {
                if map[self.y][self.x] != Location::Corridor {
                    destinations.push((x, 1, steps_made));
                }
            } else if map[2][x] == Location::Room(self.species) {
                // Try moving into the room
                let mut steps_made_entering_the_room = 0;
                let mut y = 2;
                let mut destination = None;
                let mut can_enter = true;
                let mut stopped = false;

                while map[y][x] != Location::Wall {
                    let occupant = amphipods.iter().find(|a| a.x == x && a.y == y);

                    if let Some(occupant) = occupant {
                        if occupant.species != self.species {
                            can_enter = false;
                            break;
                        }
                        stopped = true;
                    } else if !stopped {
                        steps_made_entering_the_room += 1;
                        destination = Some(y)
                    }

                    y += 1;
                }

                if can_enter {
                    if let Some(destination) = destination {
                        destinations.push((
                            x,
                            destination,
                            steps_made + steps_made_entering_the_room,
                        ));
                    }
                }
            }

            x = (x as i64 + direction) as usize;
        }

        destinations
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
struct State<const L: usize> {
    amphipods: [Amphipod; L],
    energy_spent: u64,
    estimated_energy_required: u64,
}

impl<const L: usize> Ord for State<L> {
    fn cmp(&self, other: &Self) -> Ordering {
        (other.energy_spent + other.estimated_energy_required)
            .cmp(&(self.energy_spent + self.estimated_energy_required))
            .then_with(|| self.amphipods.cmp(&other.amphipods))
    }
}

impl<const L: usize> PartialOrd for State<L> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<const L: usize> State<L> {
    fn new(amphipods: [Amphipod; L], map: &Vec<Vec<Location>>, energy_spent: u64) -> Self {
        let mut total_energy = 0;
        let mut already_in_place = HashMap::new();

        // Energy to reach correct threshold
        for amphipod in amphipods {
            if let Location::Room(room) = map[amphipod.y][amphipod.x] {
                if room == amphipod.species {
                    *already_in_place.entry(amphipod.species).or_insert(0) += 1;
                    continue
                }
            }

            total_energy += match amphipod.species {
                Species::Amber => ((amphipod.x as i64 - 3).abs() + (amphipod.y as i64 - 2).abs()),
                Species::Bronze => {
                    ((amphipod.x as i64 - 5).abs() + (amphipod.y as i64 - 1).abs()) * 10
                }
                Species::Copper => {
                    ((amphipod.x as i64 - 7).abs() + (amphipod.y as i64 - 1).abs()) * 100
                }
                Species::Desert => {
                    ((amphipod.x as i64 - 9).abs() + (amphipod.y as i64 - 1).abs()) * 1000
                }
            }
        }

        // Energy to move all amphipods from threshold into room
        total_energy += Self::energy_to_fill_room(L, &already_in_place, Species::Amber);
        total_energy += Self::energy_to_fill_room(L, &already_in_place, Species::Bronze);
        total_energy += Self::energy_to_fill_room(L, &already_in_place, Species::Copper);
        total_energy += Self::energy_to_fill_room(L, &already_in_place, Species::Desert);

        Self {
            amphipods,
            energy_spent,
            estimated_energy_required: total_energy as u64,
        }
    }

    fn energy_to_fill_room(population: usize, already_in_place: &HashMap<Species, i64>, species: Species) -> i64 {
        let to_fill = population as i64 / 4 - already_in_place.get(&species).unwrap_or(&0);
        (to_fill.pow(2) + to_fill) * species.energy_consumption() as i64 / 2
    }

    fn is_organised(&self, map: &[Vec<Location>]) -> bool {
        self.amphipods
            .iter()
            .all(|a| map[a.y][a.x] == Location::Room(a.species))
    }
}

struct Burrow<const L: usize> {
    amphipods: [Amphipod; L],
    map: Vec<Vec<Location>>,
}

impl<const L: usize> From<&str> for Burrow<L> {
    fn from(diagram: &str) -> Self {
        let mut amphipods = Vec::new();

        fn x_to_room(x: usize) -> Option<Location> {
            match x {
                3 => Some(Location::Room(Species::Amber)),
                5 => Some(Location::Room(Species::Bronze)),
                7 => Some(Location::Room(Species::Copper)),
                9 => Some(Location::Room(Species::Desert)),
                _ => None,
            }
        }

        let map = diagram
            .lines()
            .enumerate()
            .map(|(y, l)| {
                l.chars()
                    .enumerate()
                    .map(|(x, c)| match c {
                        '#' | ' ' => Location::Wall,
                        '.' => {
                            if x_to_room(x).is_none() {
                                Location::Corridor
                            } else {
                                Location::Threshold
                            }
                        }
                        'A' => {
                            amphipods.push(Amphipod {
                                species: Species::Amber,
                                x,
                                y,
                            });
                            x_to_room(x).unwrap()
                        }
                        'B' => {
                            amphipods.push(Amphipod {
                                species: Species::Bronze,
                                x,
                                y,
                            });
                            x_to_room(x).unwrap()
                        }
                        'C' => {
                            amphipods.push(Amphipod {
                                species: Species::Copper,
                                x,
                                y,
                            });
                            x_to_room(x).unwrap()
                        }
                        'D' => {
                            amphipods.push(Amphipod {
                                species: Species::Desert,
                                x,
                                y,
                            });
                            x_to_room(x).unwrap()
                        }
                        _ => panic!("Invalid input!"),
                    })
                    .collect()
            })
            .collect();

        Self {
            amphipods: amphipods.try_into().unwrap(),
            map,
        }
    }
}

impl<const L: usize> Burrow<L> {
    fn energy_to_organise(&self) -> Option<u64> {
        let mut energy_spent = HashMap::<[Amphipod; L], u64>::new();
        energy_spent.insert(self.amphipods, 0);

        let mut queue = BinaryHeap::new();
        queue.push(State::new(self.amphipods, &self.map, 0));

        while let Some(state) = queue.pop() {
            if state.is_organised(&self.map) {
                return Some(state.energy_spent);
            }

            if state.energy_spent > energy_spent[&state.amphipods] {
                continue;
            }

            for i in 0..L {
                let destinations = state.amphipods[i].possible_destinations(state.amphipods, &self.map);
                for destination in destinations {
                    let mut amphipods = state.amphipods;
                    amphipods[i].x = destination.0;
                    amphipods[i].y = destination.1;

                    let candidate = State::new(
                        amphipods,
                        &self.map,
                        state.energy_spent + destination.2 * amphipods[i].species.energy_consumption(),
                    );

                    if candidate.energy_spent
                        < *energy_spent
                            .get(&candidate.amphipods)
                            .unwrap_or(&u64::MAX)
                    {
                        queue.push(candidate);
                        energy_spent.insert(candidate.amphipods, candidate.energy_spent);
                    }
                }
            }
        }

        None
    }
}