From 7debbf7acb3419bd5a16a0c4af1a13e04dc13c35 Mon Sep 17 00:00:00 2001
From: Matthew Gretton-Dann <matt@gretton-dann.org.uk>
Date: Sat, 21 Dec 2024 16:18:45 +0000
Subject: [PATCH] Code tidy up for 2024 day 21.

---
 bin/day2421.ml | 158 +++++++++++++++++++------------------------------
 lib/aoc.ml     |   2 +
 lib/aoc.mli    |   4 ++
 3 files changed, 68 insertions(+), 96 deletions(-)

diff --git a/bin/day2421.ml b/bin/day2421.ml
index 0e4344b..cd32328 100644
--- a/bin/day2421.ml
+++ b/bin/day2421.ml
@@ -1,3 +1,15 @@
+(** A Pair of characters *)
+module CharPair = struct
+  type t = char * char
+
+  let compare (x, y) (x', y') =
+    match compare y y' with 0 -> compare x x' | c -> c
+end
+
+module CharPairMap = Map.Make (CharPair)
+
+(** [pos_of_numeric_grid c] returns the [(x, y)] position of [c] in the numeric
+    grid. *)
 let pos_of_numeric_grid c =
   match c with
   | '7' -> (0, 0)
@@ -13,7 +25,11 @@ let pos_of_numeric_grid c =
   | 'A' -> (2, 3)
   | _ -> raise (invalid_arg "pos_of_numeric_grid")
 
+(** [pos_of_numeric_grid c] returns the [(x, y)] position of [c] in the
+    direction grid. *)
 let pos_of_dir_grid c =
+  (* Implementation note: We chose 'A' to have the same position in both grids
+     so that there is only one location for the hole. *)
   match c with
   | '^' -> (1, 3)
   | 'A' -> (2, 3)
@@ -22,9 +38,14 @@ let pos_of_dir_grid c =
   | '>' -> (2, 4)
   | _ -> raise (invalid_arg "pos_of_dir_grid")
 
+(** Location of the hole which the robot can not go to. *)
 let invalid_x, invalid_y = (0, 3)
 
-let find_shortest (sx, sy) (fx, fy) =
+(** [find_paths start finish] returns a list of paths (using the direction
+    keypad to get from [start] to [finish] positions.
+
+    The routing is picked to avoid the invalid location. *)
+let find_paths (sx, sy) (fx, fy) =
   let b = Buffer.create 6 in
   let result = [] in
   let result =
@@ -51,31 +72,13 @@ let find_shortest (sx, sy) (fx, fy) =
     end
     else result
   in
-  List.sort_uniq compare result
-
-let[@warning "-32"] print_list lst =
-  let rec impl = function
-    | [] -> ()
-    | h :: t ->
-        Printf.printf "; %s" h;
-        impl t
-  in
-  match lst with
-  | [] -> print_endline "[]"
-  | h :: t ->
-      Printf.printf "[%s" h;
-      impl t;
-      Printf.printf "]\n"
-
-module CharPair = struct
-  type t = char * char
-
-  let compare (x, y) (x', y') =
-    match compare y y' with 0 -> compare x x' | c -> c
-end
-
-module CharPairMap = Map.Make (CharPair)
+  result
 
+(** [routes pos_of_grid locs] returns a map of [(start, finish)] pairs mapping
+    to a list of paths for getting to that route. [locs] are the locations on
+    the grid to investiagte. [pos_of_grid] gives the location of each of the
+    [locs]. The returned map contains routes from each element in [locs] to
+    every element. *)
 let routes pos_of_grid locs =
   let rec impl acc pos =
     let start = pos_of_grid pos in
@@ -83,7 +86,7 @@ let routes pos_of_grid locs =
     | [] -> acc
     | h :: t ->
         impl
-          (CharPairMap.add (pos, h) (find_shortest start (pos_of_grid h)) acc)
+          (CharPairMap.add (pos, h) (find_paths start (pos_of_grid h)) acc)
           pos t
   in
   let rec impl' acc = function
@@ -92,60 +95,20 @@ let routes pos_of_grid locs =
   in
   impl' CharPairMap.empty locs
 
+(** Elements on the number grid *)
 let num_grid = [ '0'; '1'; '2'; '3'; '4'; '5'; '6'; '7'; '8'; '9'; 'A' ]
+
+(** Routes between positions on the number grid *)
 let num_routes = routes pos_of_numeric_grid num_grid
+
+(** Elements on the direction grid *)
 let dir_grid = [ '<'; '>'; 'v'; '^'; 'A' ]
+
+(** Routes between positions on the direction grid. *)
 let dir_routes = routes pos_of_dir_grid dir_grid
 
-let cross_product lst lst' =
-  let rec impl' acc app' = function
-    | [] -> acc
-    | h :: t -> impl' ((h ^ app') :: acc) app' t
-  in
-  let rec impl acc = function
-    | [] -> acc
-    | h :: t -> impl (impl' acc h lst) t
-  in
-  if List.is_empty lst then lst'
-  else if List.is_empty lst' then lst
-  else impl [] lst'
-
-let find_route route_map presses =
-  let rec impl acc from presses =
-    match Seq.uncons presses with
-    | None -> acc
-    | Some (h, t) ->
-        let routes = CharPairMap.find (from, h) route_map in
-        impl (cross_product acc routes) h t
-  in
-  impl [] 'A' (String.to_seq presses)
-
-let find_routes route_map press_list =
-  let rec impl acc = function
-    | [] -> acc
-    | h :: t -> impl (acc @ find_route route_map h) t
-  in
-  impl [] press_list
-
-let min_length = List.fold_left (fun acc x -> min acc (String.length x)) max_int
-
-let human_insns2 count code =
-  let rec robot_round codes n =
-    if n >= count then codes
-    else robot_round (find_routes dir_routes codes) (n + 1)
-  in
-  let first_robot_insns = find_routes num_routes [ code ] in
-  let human = robot_round first_robot_insns ~-1 in
-  min_length human
-
-let complexity2 count code =
-  let len = human_insns2 count code in
-  let num = int_of_string (String.sub code 0 (String.length code - 1)) in
-  num * len
-
-let part1a count codes =
-  List.map (complexity2 count) codes |> List.fold_left ( + ) 0
-
+(** [initial_cost_map grid] returns a map for the initial costs (1) of moving
+    between different positions on [grid]. *)
 let initial_cost_map grid =
   let rec impl' acc f = function
     | [] -> acc
@@ -157,7 +120,9 @@ let initial_cost_map grid =
   in
   impl CharPairMap.empty grid
 
-let calc_costs cost_map steps =
+(** [calc_cost cost_map steps] calculates the cost of following [steps].
+    [cost_map] gives the cost of moving between each position. *)
+let calc_cost cost_map steps =
   let rec impl acc from seq =
     match Seq.uncons seq with
     | None -> acc
@@ -165,37 +130,38 @@ let calc_costs cost_map steps =
   in
   impl 0 'A' (String.to_seq steps)
 
+(** [get_next_level_costs route_map cost_map] gets the cost map which
+    corresponds to [route_map] with costs [cost_map]. *)
 let get_next_level_costs route_map cost_map =
   let impl routes =
-    List.fold_left (fun acc x -> min acc (calc_costs cost_map x)) max_int routes
+    List.map (calc_cost cost_map) routes
+    |> List.fold_left (fun acc x -> min acc x) max_int
   in
   CharPairMap.map impl route_map
 
-let human_insns3 count code =
-  let first_round = initial_cost_map dir_grid in
-  let rec robot_round cost_map n =
-    if n >= count then cost_map
-    else robot_round (get_next_level_costs dir_routes cost_map) (n + 1)
+(** [min_code_cost count code] returns the number of buttons a human needs to
+    press to get [code] entered when indirected through [count] robots. *)
+let min_code_cost count code =
+  let number_costs =
+    Aoc.apply_n count
+      (get_next_level_costs dir_routes)
+      (initial_cost_map dir_grid)
   in
-  let number_costs = robot_round first_round 0 in
   let number_costs = get_next_level_costs num_routes number_costs in
-  let human = calc_costs number_costs code in
-  human
+  calc_cost number_costs code
 
-let complexity3 count code =
-  let len = human_insns3 count code in
+(** [get_code_complexity count code] returns the complexity of a given code when
+    there are [count] robots involved. *)
+let get_code_complexity count code =
+  let len = min_code_cost count code in
   let num = int_of_string (String.sub code 0 (String.length code - 1)) in
   num * len
 
-let part1b count codes =
-  List.map (complexity3 count) codes |> List.fold_left ( + ) 0
+(** [part count codes] returns the puzzle rest when there are [count] robots,
+    and you need to enter [codes]. *)
+let part count codes =
+  List.map (get_code_complexity count) codes |> List.fold_left ( + ) 0
 
 let _ =
   Aoc.main Aoc.strings_of_file
-    [
-      (string_of_int, part1a 1);
-      (string_of_int, part1b 2);
-      (string_of_int, part1b 25);
-    ]
-
-(* too high 524110008179112 *)
+    [ (string_of_int, part 2); (string_of_int, part 25) ]
diff --git a/lib/aoc.ml b/lib/aoc.ml
index 35aa405..53d8d97 100644
--- a/lib/aoc.ml
+++ b/lib/aoc.ml
@@ -99,3 +99,5 @@ let memoize memo f value =
       let x = f value in
       Hashtbl.add memo value x;
       x
+
+let rec apply_n n fn arg = if n <= 0 then arg else apply_n (n - 1) fn (fn arg)
diff --git a/lib/aoc.mli b/lib/aoc.mli
index c9682a8..678e5ad 100644
--- a/lib/aoc.mli
+++ b/lib/aoc.mli
@@ -35,6 +35,10 @@ val memoize : ('a, 'b) Hashtbl.t -> ('a -> 'b) -> 'a -> 'b
     is used to cache results, so repeated calls with the same [value] will not
     call [f] again. *)
 
+val apply_n : int -> ('a -> 'a) -> 'a -> 'a
+(** [apply_n n fn arg] is equivalent to [(fn (fn ... (fn (fn arg))))] where [fn]
+    is called [n] times.*)
+
 (** Module representing a pair of integers, useful for Set.Make *)
 module IntPair : sig
   type t = int * int