Tidy up 2024 day 20.

bin/day2420.ml

@@ -1,7 +1,5 @@
-module IntMap = Map.Make (Int)
-
-let update_value = function None -> Some 1 | Some x -> Some (x + 1)
-
+(** [populate_grid grid start] does a depth-first search through [grid] to find
+    the route from [start] to the end. *)
 let populate_grid grid start =
   let costs = Array.make (Aoc.Grid.length grid) max_int in
   let rec step acc cost = function
@@ -28,11 +26,16 @@ let populate_grid grid start =
   let costs = dfs 0 [ start ] in
   costs
 
-let manhattan_distance (x, y) (x', y') = abs (x - x') + abs (y - y')
+(** [manhattan_distance2 p p'] returns the Manhattan distance between two points
+    on a 2-D plane. *)
+let manhattan_distance2 (x, y) (x', y') = abs (x - x') + abs (y - y')
 
+(** [within_distance pos distance] returns all points that are at most
+    [distance] units away from [pos] when measured using the Manhattan distance.
+*)
 let within_distance (x, y) distance =
   let rec impl' acc y' x' =
-    if manhattan_distance (x, y) (x', y') > distance then acc
+    if manhattan_distance2 (x, y) (x', y') > distance then acc
     else impl' ((x', y') :: acc) y' (x' + 1)
   in
   let rec impl acc y' =
@@ -41,43 +44,42 @@ let within_distance (x, y) distance =
   in
   impl [] (y - distance)
 
-let find_cost2 map depth_first grid idx length =
+(** [find_cost min_amount depth_first grid length idx] returns the number of
+    cheat routes starting at [idx] which have a saving of at least [min_amount]
+    and are no longer than [length]. [depth_first] is the cost map, and [grid]
+    is the grid. *)
+let find_cost min_amount depth_first grid length idx =
+  (* because there is only one route through the grid we can specialize and look
+     to see how much we can save by going from [idx] to any other grid position
+     within [length] units (manhattan distance).  The saving is the cost of
+     gettimg to idx' from idx via the old route - the cost via the new route. *)
   let saving idx' =
     let cost = depth_first.(idx) in
     let cost' = depth_first.(idx') in
-    let saving =
     cost' - cost
-      - manhattan_distance
+    - manhattan_distance2
         (Aoc.Grid.pos_of_idx grid idx)
         (Aoc.Grid.pos_of_idx grid idx')
   in
-    saving
-  in
-  let rec impl acc = function
-    | [] -> acc
-    | h :: t -> impl (IntMap.update (saving h) update_value acc) t
-  in
   within_distance (Aoc.Grid.pos_of_idx grid idx) length
   |> List.filter (Aoc.Grid.pos_is_valid grid)
   |> List.map (Aoc.Grid.idx_of_pos grid)
-  |> List.filter (fun x -> Aoc.Grid.get_by_idx grid x <> '#')
-  |> List.filter (fun x -> depth_first.(x) - depth_first.(idx) >= 0)
-  |> impl map
+  |> List.filter (fun idx' -> depth_first.(idx') <> max_int)
+  |> List.filter (fun idx' -> depth_first.(idx') - depth_first.(idx) >= 0)
+  |> List.map saving
+  |> List.filter (( <= ) min_amount)
+  |> List.length
 
-let find_cost_reductions2 grid start cheat_length =
-  let depth_first = populate_grid grid start in
-  let rec impl acc idx =
-    if idx >= Aoc.Grid.length grid then acc
-    else
-      let acc = find_cost2 acc depth_first grid idx cheat_length in
-      impl acc (idx + 1)
-  in
-  impl IntMap.empty 0
-
-let part2 cheat_length (grid, start) =
-  let map = find_cost_reductions2 grid start cheat_length in
-  let map = IntMap.filter (fun k _ -> k >= 100) map in
-  IntMap.fold (fun _ v a -> v + a) map 0
+(** [find_cost_reductions min_amount cheat_length (grid, start)] returns the
+    number of cheat-routes that can be found in [grid] starting at [start] that
+    save at least [min_amount] moves and are no longer than [cheat_length]
+    units. *)
+let find_cost_reductions min_amount cheat_length (grid, start) =
+  let costs = populate_grid grid start in
+  Seq.ints 0
+  |> Seq.take (Aoc.Grid.length grid)
+  |> Seq.map (find_cost min_amount costs grid cheat_length)
+  |> Seq.fold_left ( + ) 0
 
 let find_start grid =
   match Aoc.Grid.idx_from_opt grid 0 'S' with
@@ -90,4 +92,8 @@ let data_of_file fname =
   (grid, start)
 
 let _ =
-  Aoc.main data_of_file [ (string_of_int, part2 2); (string_of_int, part2 20) ]
+  Aoc.main data_of_file
+    [
+      (string_of_int, find_cost_reductions 100 2);
+      (string_of_int, find_cost_reductions 100 20);
+    ]