Comment 2024 day 16.

bin/day2416.ml

@@ -1,16 +1,35 @@
+(** [find grid c] returns the position of the character [c] in [grid]. *)
 let find grid c =
   match Aoc.Grid.idx_from_opt grid 0 c with
   | None -> failwith "find"
   | Some idx -> Aoc.Grid.pos_of_idx grid idx
 
+(** [input_of_file fname] returns a [(grid, start_pos)] pair parsed from
+    [fname]. *)
 let input_of_file fname =
   let grid = Aoc.Grid.of_file fname in
   let start_pos = find grid 'S' in
   (grid, start_pos)
 
-let rec dijkstra visit check_end states =
+(** [dijkstra visit check_end states] executes Dijkstra's algorithm.
+
+    [visit cost state] is called to visit [state] with [cost]. It should mark
+    [state] as visited, and return a list of [(cost, state)] pairs which contain
+    new states to examine. The returned list should be sorted by [cost].
+
+    [check_end state] should return [true] if and only if [state] is an end
+    state.
+
+    [states] is a list of [(cost, state)] pairs ordered by [cost].
+
+    [dijkstra] returns [None] if no path is found to the destination. It returns
+    [Some (cost, state, remaining_states)] if a route is found. [cost] is the
+    cost of getting to [state]. [remaining_states] is a list of the remaining
+    states which can be passed back to [dijkstra] if we want to find further
+    paths. *)
+let rec dijkstra visit check_end =
   let compare_costs (lhs, _) (rhs, _) = compare lhs rhs in
-  match states with
+  function
   | [] -> None
   | (cost, state) :: t ->
       if check_end state then Some (cost, state, t)
@@ -18,6 +37,8 @@ let rec dijkstra visit check_end states =
         let new_states = visit cost state |> List.merge compare_costs t in
         dijkstra visit check_end new_states
 
+(** [visited_idx grid state] returns the index into the visited array for [grid]
+    at a given [state]. *)
 let visited_idx grid ((dx, dy), p) =
   let add =
     match (dx, dy) with
@@ -29,6 +50,9 @@ let visited_idx grid ((dx, dy), p) =
   in
   (Aoc.Grid.idx_of_pos grid p * 4) + add
 
+(** [visit grid visited cost state] visits [state] with [cost] in [grid].
+    [visited] is an array of visited states, and is updated as we visit. It
+    returns a list of new [(cost, state)] pairs to visit. *)
 let visit grid visited_grid cost state =
   let (dx, dy), ((x, y) as p) = List.hd state in
   let has_visited = visited_grid.(visited_idx grid (List.hd state)) in
@@ -42,16 +66,19 @@ let visit grid visited_grid cost state =
       (cost + 1000, ((dy, -dx), p) :: state);
     ])
 
+(** [has_visited grid visited_grid (cost, state)] returns true if we have
+    visited the [state]. *)
 let has_visited grid visited_grid (cost, state) =
   visited_grid.(visited_idx grid (List.hd state)) < cost
 
+(** [visit grid visited cost state] visits [state] with [cost] in [grid].
+    [visited] is an array of visited states, and is updated as we visit. It
+    returns a list of new [(cost, state)] pairs to visit. *)
 let visit_max grid visited_grid cost state =
   let (dx, dy), ((x, y) as p) = List.hd state in
   if has_visited grid visited_grid (cost, state) then []
   else if Aoc.Grid.get_by_pos grid p = '#' then []
   else (
-    Printf.printf "%d (%d, %d)\n" cost x y;
-    flush stdout;
     visited_grid.(visited_idx grid (List.hd state)) <- cost;
     [
       (cost + 1, ((dx, dy), (x + dx, y + dy)) :: state);
@@ -60,14 +87,16 @@ let visit_max grid visited_grid cost state =
     ]
     |> List.filter (fun x -> not (has_visited grid visited_grid x)))
 
-let[@warning "-32"] print_point (x, y) = Printf.printf "(%d, %d)" x y
+(** [check_end grid state] returns [true] if [state] is at the end location in
-
+    [grid]. *)
 let check_end grid state =
   let _, p = List.hd state in
   Aoc.Grid.get_by_pos grid p = 'E'
 
-let check_end2 grid _ state = check_end grid state
+(** [part1 (grid, start_pos)] returns solution to part 1.
 
+    This part does a simple Dijkstra algorithm over the grid finding the
+    shortest path possible. *)
 let part1 (grid, start_pos) =
   let visited_grid = Array.make (Aoc.Grid.length grid * 4) false in
   match
@@ -77,13 +106,18 @@ let part1 (grid, start_pos) =
   | None -> failwith "part"
   | Some (cost, _, _) -> cost
 
+(** [part2 (grid, start_pos)] returns solution to part 2.
+
+    We reuse part 1 to find the cost of getting to the exit. Then we redo the
+    Dijkstra algorithm to find all walks with that cost.
+
+    Finally we get all the positions visited, de-duplicate, and count the length
+    of the list. This produces the number of locations for benches. *)
 let part2 (grid, start_pos) =
   let cost = part1 (grid, start_pos) in
   let visited_grid = Array.make (Aoc.Grid.length grid * 4) cost in
   let rec impl acc lst =
-    match
+    match dijkstra (visit_max grid visited_grid) (check_end grid) lst with
-      dijkstra (visit_max grid visited_grid) (check_end2 grid visited_grid) lst
-    with
     | None -> acc
     | Some (_, states, remainder) ->
         List.iter (fun x -> visited_grid.(visited_idx grid x) <- 0) states;