d21

21/main.ml

@@ -0,0 +1,132 @@
+open Printf;;
+
+let rec list_of_lines in_file =
+  try
+    let line = input_line in_file in
+    line :: list_of_lines(in_file)
+  with End_of_file ->
+    close_in in_file;
+    []
+
+let is_valid_position_numpad (x, y) =
+  x >= 0 && x < 3 && y >= 0 && y < 4 && not (x = 0 && y = 3)
+
+let is_valid_position_keypad (x, y) =
+  x >= 0 && x < 3 && y >= 0 && y < 2 && not (x = 0 && y = 0)
+
+let get_keypad_coords symbol = 
+  match symbol with
+  | '<' -> (0, 1)
+  | '>' -> (2, 1)
+  | '^' -> (1, 0)
+  | 'v' -> (1, 1)
+  | 'A' -> (2, 0)
+  | _ -> raise (Invalid_argument "Invalid symbol")
+  
+let get_numpad_coords symbol = 
+  if symbol = 'A' then
+    (2, 3)
+  else
+    let n = (Char.code symbol) - (Char.code '0') in
+    if n = 0 then
+      (1, 3)
+    else
+      ((n - 1) mod 3, (9 - n) / 3)
+
+let get_dir_vector dir_symbol =
+  match dir_symbol with
+  | '<' -> (-1, 0)
+  | '>' -> (1, 0)
+  | '^' -> (0, -1)
+  | 'v' -> (0, 1)
+  | _ -> raise (Invalid_argument "Invalid direction")
+
+let allowed_moves prev =
+  match prev with
+  | Some 'A' | None -> ['<'; '^'; '>'; 'v']
+  | Some '<' -> ['<'; '^'; 'v']
+  | Some '>' -> ['^'; '>'; 'v']
+  | Some '^' -> ['<'; '^'; '>']
+  | Some 'v' -> ['<'; '>'; 'v']
+  | _ -> raise (Invalid_argument "Invalid move")
+
+module CombinationSearch =
+struct
+  type t = int * (int * (int * int) * char * char list)
+  let compare  a b =
+    Stdlib.compare a b
+end
+module CombinationSearchSet = Set.Make(CombinationSearch)
+
+let rec find_best_combination_dijkstra lower_level_cost_fun is_valid_fun to_pos queue =
+  match CombinationSearchSet.min_elt_opt queue with
+  | None -> raise Not_found
+  | Some (score, (partial_score, pos, prev_dir, history)) ->
+    if pos = to_pos then begin
+      score
+    end
+    else
+      let queue_removed = CombinationSearchSet.remove (score, (partial_score, pos, prev_dir, history)) queue in
+      Some prev_dir
+      |> allowed_moves
+      |> List.filter_map (fun dir ->
+        let vec = get_dir_vector dir in
+        let new_pos = (fst pos + fst vec, snd pos + snd vec) in
+        let lower_cost = lower_level_cost_fun prev_dir dir in
+        let finish_cost = lower_level_cost_fun dir 'A' in
+        let new_partial_score = partial_score + lower_cost in
+        if is_valid_fun new_pos then
+          Some (new_partial_score + finish_cost, (new_partial_score, new_pos, dir, dir :: history))
+        else
+          None)
+      |> List.fold_left (fun acc state -> CombinationSearchSet.add state acc) queue_removed
+      |> find_best_combination_dijkstra lower_level_cost_fun is_valid_fun to_pos
+
+let find_best_combination lower_level_cost_fun is_valid_fun get_coords_fun from_symbol to_symbol =
+  let coords1 = get_coords_fun from_symbol in
+  let coords2 = get_coords_fun to_symbol in
+  find_best_combination_dijkstra lower_level_cost_fun is_valid_fun coords2 ([(1, (0, coords1, 'A', []))] |> CombinationSearchSet.of_list)
+
+let get_numerical_value str =
+  int_of_string (String.sub str 0 (String.length str - 1))
+
+let rec get_keypad_combination_length cache depth from_symbol to_symbol =
+  match Hashtbl.find_opt cache (from_symbol, to_symbol, depth) with
+  | Some score -> score
+  | None ->
+    let best = find_best_combination (get_keypad_combination_length cache (depth - 1))
+                  is_valid_position_keypad get_keypad_coords from_symbol to_symbol in
+    Hashtbl.add cache (from_symbol, to_symbol, depth) best;
+    best
+
+let get_numpad_combination_length cache keypad_depth from_symbol to_symbol = 
+  let result = find_best_combination (get_keypad_combination_length cache keypad_depth)
+      is_valid_position_numpad get_numpad_coords from_symbol to_symbol in
+  result
+
+let get_numpad_str_combination_length cache keypad_depth str =
+  str
+  |> String.fold_left (fun (cost, prev) c -> (cost + get_numpad_combination_length cache keypad_depth prev c, c)) (0, 'A')
+  |> fst
+
+let sum_solutions cache keypad_depth lst =
+  lst
+  |> List.map (fun str -> (get_numpad_str_combination_length cache keypad_depth str, get_numerical_value str))
+  |> List.map (fun (solution, value) -> solution * value)
+  |> List.fold_left Int.add 0
+
+let init_cache =
+  let cache = Hashtbl.create 256 in
+  let key_options = ['<'; '>'; '^'; 'v'; 'A'] in
+  key_options |> List.iter (fun c1 ->
+    key_options |> List.iter (fun c2 -> Hashtbl.add cache (c1, c2, 0) 1));
+  cache
+
+let () =
+  let () = Printexc.record_backtrace true in
+  let f = open_in "input.txt" in
+  let lines = list_of_lines f in
+  let cache = init_cache in
+  let result = sum_solutions cache 2 lines in
+  let result2 = sum_solutions cache 25 lines in
+  printf "%d\n%d\n" result result2