Comment and tidy up 2024 day 18

bin/day2418.ml

@@ -1,12 +1,8 @@
+(** [pairs_of_ints lst] returns a pair from a list of two elements. *)
 let pairs_of_ints = function
   | [ h; h' ] -> (h, h')
   | _ -> raise (Invalid_argument "pairs_of_ints")
 
-let pairs_of_file fname =
-  Aoc.strings_of_file fname
-  |> List.map (Aoc.ints_of_string ~sep:",")
-  |> List.map pairs_of_ints
-
 (** [dijkstra visit check_end states] executes Dijkstra's algorithm.
 
     [visit cost state] is called to visit [state] with [cost]. It should mark
@@ -26,68 +22,112 @@ let pairs_of_file fname =
 let rec dijkstra visit check_end =
   let compare_costs (lhs, _) (rhs, _) = compare lhs rhs in
   function
-  | [] -> failwith "dijkstra"
+  | [] -> None
   | (cost, state) :: t ->
-      if check_end state then (cost, state)
+      if check_end state then Some (cost, state)
       else
         let new_states = visit cost state |> List.merge compare_costs t in
         dijkstra visit check_end new_states
 
-type grid = { grid : bool array; width : int }
+type 'a grid = { grid : 'a array; width : int }
 
+(** [grid_is_valid_pos grid (x, y)] returns true if (x, y) is a valid position
+*)
 let grid_is_valid_pos grid (x, y) =
   x >= 0 && x < grid.width && y >= 0 && y < grid.width
 
+(** Get the index into the grid from an x, y position. *)
 let grid_idx_by_pos grid (x, y) = x + (y * grid.width)
 
+(** Set the value of the position (x, y) to v in grid. *)
 let grid_set_by_pos grid p v =
   assert (grid_is_valid_pos grid p);
   let idx = grid_idx_by_pos grid p in
   grid.grid.(idx) <- v
 
+(** Get the value of the position (x, y) in grid. *)
 let grid_get_by_pos grid p =
   assert (grid_is_valid_pos grid p);
   let idx = grid_idx_by_pos grid p in
   grid.grid.(idx)
 
-let make_grid width rocks =
-  let grid = { grid = Array.make (width * width) false; width } in
-  let rec impl = function
-    | [] -> grid
-    | p :: t ->
-        grid_set_by_pos grid p true;
-        impl t
-  in
-  impl rocks
+(** [grid_of_rocks w rocks] returns a [w * w] grid with [grid.(x + y * w)]
+    indicating whether the space is empty ([=max_int]) or which rock it is (0
+    based). *)
+let grid_of_rocks width rocks =
+  let grid = { grid = Array.make (width * width) max_int; width } in
+  let add_rock idx p = grid_set_by_pos grid p idx in
+  List.iteri add_rock rocks;
+  grid
 
-let visit grid has_visited cost state =
+(** [visit grid has_visited count cost pos] visits the location pos marking it
+    as visited and returning a list of [(cost, pos)] pairs of next locations to
+    examine.
+
+    [grid] is the grid of rocks, [has_visited] is an array of bools indicating
+    whether a position has already been visited, and [count] is how many rocks
+    have fallen. *)
+let visit grid has_visited count cost state =
   if not (grid_is_valid_pos grid state) then []
   else if has_visited.(grid_idx_by_pos grid state) then []
-  else if grid_get_by_pos grid state then []
+  else if grid_get_by_pos grid state < count then []
   else
     let x, y = state in
     has_visited.(grid_idx_by_pos grid state) <- true;
     [
-      (cost + 1, (x - 1, y));
       (cost + 1, (x + 1, y));
-      (cost + 1, (x, y - 1));
+      (cost + 1, (x - 1, y));
       (cost + 1, (x, y + 1));
+      (cost + 1, (x, y - 1));
     ]
 
-let check_end dest state = dest = state
+(** [grid_of_file w fname] returns a grid of width & height [w] populated with
+    rocks described in the file [fname]. *)
+let grid_of_file width fname =
+  Aoc.strings_of_file fname
+  |> List.map (Aoc.ints_of_string ~sep:",")
+  |> List.map pairs_of_ints |> grid_of_rocks width
 
-let part1 lst =
-  let width = 71 in
-  let count = 1024 in
-  let grid =
-    List.to_seq lst |> Seq.take count |> List.of_seq |> make_grid width
-  in
-  let has_visited = Array.make (width * width) false in
-  let cost, _ =
-    dijkstra (visit grid has_visited)
-      (check_end (width - 1, width - 1))
+(** [find_route_length count grid] calculates the route from the top-left
+    position in [grid] to the bottom right if [count] rocks have fallen. It
+    returns [None] if no route is possible or [Some (cost, pos)] if the route is
+    possible. *)
+let find_route_length count grid =
+  let has_visited = Array.make (Array.length grid.grid) false in
+  dijkstra
+    (visit grid has_visited count)
+    (( = ) (grid.width - 1, grid.width - 1))
     [ (0, (0, 0)) ]
-  in
-  cost
 
-let _ = Aoc.main pairs_of_file [ (string_of_int, part1) ]
+(** [part1 count rocks] returns how long it takes to navigate the grid [rocks]
+    when [count] rocks have fallen. *)
+let part1 count rocks =
+  match find_route_length count rocks with
+  | None -> failwith "part1"
+  | Some (cost, _) -> cost
+
+(** [part2 start_count grid] returns the location of the first rock to fall into
+    [grid] which makes it impossible to get from the top-left to bottom-right.
+*)
+let part2 start_count grid =
+  let rec impl count =
+    match find_route_length count grid with
+    | None -> count
+    | Some _ -> impl (count + 1)
+  in
+  let count = impl start_count in
+  match Array.find_index (( = ) (count - 1)) grid.grid with
+  | None -> failwith "part2"
+  | Some idx -> idx
+
+(** [string_of_idx width idx] prints the (x, y) location for a given index in a
+    grid. *)
+let string_of_idx width idx =
+  Printf.sprintf "%d,%d" (idx mod width) (idx / width)
+
+(** Width of grid *)
+let width = 71
+
+let _ =
+  Aoc.main (grid_of_file width)
+    [ (string_of_int, part1 1024); (string_of_idx width, part2 1024) ]