2024 day 16 part 2

Quite slow takes ~11 mins to run.

bin/day2416.ml

@@ -0,0 +1,100 @@
+let find grid c =
+  match Aoc.Grid.idx_from_opt grid 0 c with
+  | None -> failwith "find"
+  | Some idx -> Aoc.Grid.pos_of_idx grid idx
+
+let input_of_file fname =
+  let grid = Aoc.Grid.of_file fname in
+  let start_pos = find grid 'S' in
+  (grid, start_pos)
+
+let rec dijkstra visit check_end states =
+  let compare_costs (lhs, _) (rhs, _) = compare lhs rhs in
+  match states with
+  | [] -> None
+  | (cost, state) :: t ->
+      if check_end state then Some (cost, state, t)
+      else
+        let new_states =
+          visit cost state |> List.sort compare_costs
+          |> List.merge compare_costs t
+        in
+        dijkstra visit check_end new_states
+
+let visited_idx grid ((dx, dy), p) =
+  let add =
+    match (dx, dy) with
+    | 1, 0 -> 0
+    | 0, 1 -> 1
+    | -1, 0 -> 2
+    | 0, -1 -> 3
+    | _ -> failwith "visited_idx"
+  in
+  (Aoc.Grid.idx_of_pos grid p * 4) + add
+
+let visit grid visited_grid cost state =
+  let (dx, dy), ((x, y) as p) = List.hd state in
+  let has_visited = visited_grid.(visited_idx grid (List.hd state)) in
+  if has_visited then []
+  else if Aoc.Grid.get_by_pos grid p = '#' then []
+  else (
+    visited_grid.(visited_idx grid (List.hd state)) <- true;
+    [
+      (cost + 1, ((dx, dy), (x + dx, y + dy)) :: state);
+      (cost + 1000, ((-dy, dx), p) :: state);
+      (cost + 1000, ((dy, -dx), p) :: state);
+    ])
+
+let visit_max grid visited_grid max_cost cost state =
+  let has_visited = visited_grid.(visited_idx grid (List.hd state)) < cost in
+  let (dx, dy), ((x, y) as p) = List.hd state in
+  if has_visited then []
+  else if cost > max_cost then []
+  else if Aoc.Grid.get_by_pos grid p = '#' then []
+  else (
+    (*Printf.printf "%d (%d, %d)\n" cost x y;
+    flush stdout;*)
+    visited_grid.(visited_idx grid (List.hd state)) <- cost;
+    [
+      (cost + 1, ((dx, dy), (x + dx, y + dy)) :: state);
+      (cost + 1000, ((-dy, dx), p) :: state);
+      (cost + 1000, ((dy, -dx), p) :: state);
+    ])
+
+let[@warning "-32"] print_point (x, y) = Printf.printf "(%d, %d)" x y
+
+let check_end grid state =
+  let _, p = List.hd state in
+  Aoc.Grid.get_by_pos grid p = 'E'
+
+let check_end2 grid _ state = check_end grid state
+
+let part1 (grid, start_pos) =
+  let visited_grid = Array.make (Aoc.Grid.length grid * 4) false in
+  match
+    dijkstra (visit grid visited_grid) (check_end grid)
+      [ (0, [ ((1, 0), start_pos) ]) ]
+  with
+  | None -> failwith "part"
+  | Some (cost, _, _) -> cost
+
+let part2 (grid, start_pos) =
+  let cost = part1 (grid, start_pos) in
+  let visited_grid = Array.make (Aoc.Grid.length grid * 4) max_int in
+  let rec impl acc lst =
+    match
+      dijkstra
+        (visit_max grid visited_grid cost)
+        (check_end2 grid visited_grid)
+        lst
+    with
+    | None -> acc
+    | Some (_, states, remainder) ->
+        List.iter (fun x -> visited_grid.(visited_idx grid x) <- 0) states;
+        impl (states :: acc) remainder
+  in
+  impl [] [ (0, [ ((1, 0), start_pos) ]) ]
+  |> List.concat |> List.map snd |> List.sort_uniq compare |> List.length
+
+let _ =
+  Aoc.main input_of_file [ (string_of_int, part1); (string_of_int, part2) ]

bin/dune

@@ -14,7 +14,8 @@
   day2412
   day2413
   day2414
-  day2415)
+  day2415
+  day2416)
  (names
   day2401
   day2402
@@ -30,5 +31,6 @@
   day2412
   day2413
   day2414
-  day2415)
+  day2415
+  day2416)
  (libraries str aoc))