Day 17: Clumsy Crucible

This is a very straightforward application of pathfinding (Dijkstra, A*, etc.), but the state transition function was a little more complicated than usual for grid pathfinding.

I wanted to experiment with a search algorithm called Anytime Focal Search (PDF), which doesn't really improve over plain A* a lot here, but there was no way of knowing that without trying it. (Edit: I just realized that the version of AFS I implemented is exactly equivalent to A* here, because I don't have a more sophisticated heuristic than the manhattan distance).

Anytime Focal Search is good when there are many paths of similar cost, or when you need an answer quickly, even if it is sub-optimal. Giving the algorithm more time to run then successively improves the path found, until it is optimal.

I'm also using a RadixHeap, a monotone priority queue useful for pathfinding problems, because that's slightly faster than the BinaryHeap that comes with the standard library.

use std::rc::Rc;
use std::{cmp::Reverse, hash::Hash};

use ahash::HashSet;
use glam::IVec2;
use radix_heap::RadixHeapMap;

use crate::{
    parser::{Grid2D, HeatLoss, PuzzleInput},
    part1::Direction,
};

// Search state
#[derive(Clone, Debug)]
pub struct State {
    grid: Rc<Grid2D<HeatLoss>>, // Grid2D is a wrapper around ndarray::Array2
    coordinate: IVec2,
    direction: Option<Direction>,
    minimum_straight_moves_remaining: u8,
    maximum_straight_moves_remaining: u8,
}

impl PartialEq for State {
    fn eq(&self, other: &Self) -> bool {
        self.coordinate == other.coordinate
            && self.direction == other.direction
            && self.minimum_straight_moves_remaining == other.minimum_straight_moves_remaining
            && self.maximum_straight_moves_remaining == other.maximum_straight_moves_remaining
    }
}

impl Eq for State {}

impl Hash for State {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.coordinate.hash(state);
        self.direction.hash(state);
        self.minimum_straight_moves_remaining.hash(state);
        self.maximum_straight_moves_remaining.hash(state);
    }
}

// States reachable from `s` and the cost of moving there from s (the heat loss incurred) 
fn successors(s: &State) -> Vec<(State, usize)> {
    let mut result: Vec<(State, usize)> = Vec::new();

    let mut add_successor = |direction: Direction| {
        let mut new_state = s.clone();

        let going_straight = s.direction == Some(direction) || s.direction.is_none();

        if going_straight && s.maximum_straight_moves_remaining == 0 {
            return;
        }

        if !going_straight && s.minimum_straight_moves_remaining > 0 {
            return;
        }

        new_state.direction = Some(direction);

        match new_state.direction {
            Some(Direction::North) => new_state.coordinate.y -= 1,
            Some(Direction::South) => new_state.coordinate.y += 1,
            Some(Direction::West) => new_state.coordinate.x -= 1,
            Some(Direction::East) => new_state.coordinate.x += 1,
            None => unreachable!("Direction should never be None"),
        }

        if going_straight {
            new_state.maximum_straight_moves_remaining -= 1;
            new_state.minimum_straight_moves_remaining =
                new_state.minimum_straight_moves_remaining.saturating_sub(1);
        } else {
            new_state.maximum_straight_moves_remaining = 10 - 1;
            new_state.minimum_straight_moves_remaining = 4 - 1;
        }

        if new_state.coordinate.x >= 0
            && new_state.coordinate.x < s.grid.width
            && new_state.coordinate.y >= 0
            && new_state.coordinate.y < s.grid.height
        {
            let x = new_state.coordinate.x;
            let y = new_state.coordinate.y;

            result.push((
                new_state,
                s.grid.data.get([y as usize, x as usize]).unwrap().0 as usize,
            ));
        }
    };

    if s.direction.is_none() {
        add_successor(Direction::South);
        add_successor(Direction::East);
    } else {
        add_successor(s.direction.unwrap());
        add_successor(s.direction.unwrap().turn_left());
        add_successor(s.direction.unwrap().turn_right());
    }

    result
}

// Are we at the goal?
fn success(s: &State) -> bool {
    s.coordinate.x == s.grid.width - 1
        && s.coordinate.y == s.grid.height - 1
        && s.minimum_straight_moves_remaining == 0
}

// Manhattan distance heuristic. I'm sure this could probably be improved, since
// we have to turn around if we're facing the wrong direction.
fn heuristic(s: &State) -> usize {
    (s.coordinate.x.abs_diff(s.grid.width - 1) + s.coordinate.y.abs_diff(s.grid.height - 1))
        as usize
}

// Pathfinding algorithm, similar to A*
pub fn anytime_focal_search<FN, IN, FS>(
    start: &State,
    mut successors: FN,
    mut success: FS,
) -> Option<usize>
where
    FN: FnMut(&State) -> IN,
    IN: IntoIterator<Item = (State, usize)>,
    FS: FnMut(&State) -> bool,
{
    let mut focal = RadixHeapMap::new();
    let mut seen = HashSet::default();
    let mut cur_bound = usize::MAX;
    let mut cur_sol = None;

    focal.push(Reverse(0), (0, start.clone()));

    loop {
        cur_bound -= 1;

        let mut cost = usize::MAX;

        while !focal.is_empty() {
            let (f, (g, n)) = focal.pop().unwrap();

            if f.0 > cur_bound {
                continue;
            }

            if success(&n) {
                cost = g;
                break;
            }

            for (s, c) in successors(&n).into_iter() {
                let f = g + c + heuristic(&s);

                if f <= cur_bound && seen.insert(s.clone()) {
                    focal.push(Reverse(f), (g + c, s));
                }
            }
        }

        if cost == usize::MAX {
            return cur_sol;
        }

        cur_sol = Some(cost);
        cur_bound = cost;
    }
}

// Just initialize the first state and call AFS
pub fn part2(input: PuzzleInput) -> isize {
    let initial_state = State {
        grid: Rc::new(input.grid),
        coordinate: IVec2::new(0, 0),
        maximum_straight_moves_remaining: 10,
        minimum_straight_moves_remaining: 4,
        direction: None,
    };

    anytime_focal_search(&initial_state, successors, success).unwrap() as isize
}

My solution was really slow at first, because I forgot to filter out duplicate states, but now it is reasonably fast:

day_17_bench  fastest       │ slowest       │ median        │ mean          │ samples │ iters
├─ parse      28.09 µs      │ 62.11 µs      │ 28.24 µs      │ 28.78 µs      │ 100     │ 100
├─ part1      37.71 ms      │ 63.05 ms      │ 44.29 ms      │ 45.14 ms      │ 100     │ 100
╰─ part2      124.9 ms      │ 142.1 ms      │ 125.5 ms      │ 125.9 ms      │ 100     │ 100  

Pathfinding search problem -- I knew this wanted something like Dijkstra's algorithm. I wanted the data in a grid to manage the directional aspects of it, but the actual graph to search is the location in the grid + the state of how you go there.

Simple implementation got the tests in part 1, but was pretty slow for the full input. However, it finished before I could implement a priority queue to speed up. :)

Picking up part 2, the different rules just dictate a different case for how neighbors are found in the grid, so
I cased both into may algorithm. I was able to get the tests for part 2, but needed the priority queue implementation before I could get a solution for the whole input.

Partly I blame how long the PQ implementation took on the way the heap is implemented in go -- I needed to write a wrapper around it, and there were bugs in the integration of it in my algorithm.

In the end, it's pretty messy, but I have spent far too long on it already.

# Advent of Code 2023
# Day 17: Clumsy Crucible

import heapq

with open("inputs/day_17.txt") as file:
    loss_map = [[int(char) for char in line.rstrip("\n")] for line in file]
    width = len(loss_map[0])
    height = len(loss_map)

class Node:
    def __init__(self, x, y, direction, consecutive, heat_loss, part):
        self.x = x
        self.y = y
        self.direction = direction
        self.consecutive = consecutive
        self.heat_loss = heat_loss
        self.part = part

    def __lt__(self, other):
        return self.heat_loss < other.heat_loss

heap = [Node(0, 0, None, 0, 0, 1), Node(0, 0, None, 0, 0, 2)]
heapq.heapify(heap)
seen = set()

end_x = width - 1
end_y = height - 1

while heap:
    current = heapq.heappop(heap)

    if current.x == end_x and current.y == end_y:
        print(f"Part {current.part}: {current.heat_loss}")

        for node in heap.copy():
            if node.part == current.part:
                heap.remove(node)

        continue

    for new_pos in (
        (current.x + 1, current.y, "right"),
        (current.x, current.y + 1, "down"),
        (current.x - 1, current.y, "left"),
        (current.x, current.y - 1, "up")
    ):
        if (
            0 <= new_pos[0] < width and 0 <= new_pos[1] < height
            and not (
                current.direction == "right" and new_pos[2] == "left"
                or current.direction == "down" and new_pos[2] == "up"
                or current.direction == "left" and new_pos[2] == "right"
                or current.direction == "up" and new_pos[2] == "down"
            )
            and not (
                current.direction == new_pos[2]
                and current.consecutive == (3, 10)[current.part - 1]
            )
        ):
            node = Node(
                *new_pos,
                (
                    current.consecutive + 1
                    if current.direction == new_pos[2] else 1
                ),
                current.heat_loss + loss_map[new_pos[1]][new_pos[0]],
                current.part
            )

            if current.part == 2 and node.consecutive == 1:
                if node.direction == "right":
                    change = (1, 0)
                elif node.direction == "down":
                    change = (0, 1)
                elif node.direction == "left":
                    change = (-1, 0)
                else:
                    change = (0, -1)

                for i in range(3):
                    node.x += change[0]
                    node.y += change[1]

                    try:
                        node.heat_loss += loss_map[node.y][node.x]
                    except IndexError:
                        pass

                node.consecutive = 4

            values = (node.x, node.y, node.direction, node.consecutive)

            if values not in seen:
                heapq.heappush(heap, node)
                seen.add(values)