Code tidy up for 2024 day 21.

bin/day2421.ml

@@ -0,0 +1,167 @@
+(** 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)
+  | '8' -> (1, 0)
+  | '9' -> (2, 0)
+  | '4' -> (0, 1)
+  | '5' -> (1, 1)
+  | '6' -> (2, 1)
+  | '1' -> (0, 2)
+  | '2' -> (1, 2)
+  | '3' -> (2, 2)
+  | '0' -> (1, 3)
+  | '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)
+  | '<' -> (0, 4)
+  | 'v' -> (1, 4)
+  | '>' -> (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)
+
+(** [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 =
+    if fx <> invalid_x || sy <> invalid_y then begin
+      Buffer.reset b;
+      if fx > sx then Buffer.add_string b (String.make (fx - sx) '>')
+      else if fx < sx then Buffer.add_string b (String.make (sx - fx) '<');
+      if fy > sy then Buffer.add_string b (String.make (fy - sy) 'v')
+      else if fy < sy then Buffer.add_string b (String.make (sy - fy) '^');
+      Buffer.add_char b 'A';
+      Buffer.contents b :: result
+    end
+    else result
+  in
+  let result =
+    if sx <> invalid_x || fy <> invalid_y then begin
+      Buffer.reset b;
+      if fy > sy then Buffer.add_string b (String.make (fy - sy) 'v')
+      else if fy < sy then Buffer.add_string b (String.make (sy - fy) '^');
+      if fx > sx then Buffer.add_string b (String.make (fx - sx) '>')
+      else if fx < sx then Buffer.add_string b (String.make (sx - fx) '<');
+      Buffer.add_char b 'A';
+      Buffer.contents b :: result
+    end
+    else result
+  in
+  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
+    function
+    | [] -> acc
+    | h :: t ->
+        impl
+          (CharPairMap.add (pos, h) (find_paths start (pos_of_grid h)) acc)
+          pos t
+  in
+  let rec impl' acc = function
+    | [] -> acc
+    | h :: t -> impl' (impl acc h locs) t
+  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
+
+(** [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
+    | h :: t -> impl' (CharPairMap.add (f, h) 1 acc) f t
+  in
+  let rec impl acc = function
+    | [] -> acc
+    | h :: t -> impl (impl' acc h grid) t
+  in
+  impl CharPairMap.empty grid
+
+(** [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
+    | Some (h, t) -> impl (acc + CharPairMap.find (from, h) cost_map) h t
+  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.map (calc_cost cost_map) routes
+    |> List.fold_left (fun acc x -> min acc x) max_int
+  in
+  CharPairMap.map impl route_map
+
+(** [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 = get_next_level_costs num_routes number_costs in
+  calc_cost number_costs code
+
+(** [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
+
+(** [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, part 2); (string_of_int, part 25) ]

bin/dune

@@ -19,7 +19,8 @@
   day2417
   day2418
   day2419
-  day2420)
+  day2420
+  day2421)
  (names
   day2401
   day2402
@@ -40,5 +41,6 @@
   day2417
   day2418
   day2419
-  day2420)
+  day2420
+  day2421)
  (libraries str aoc))

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)

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