refactor generalized djikstra out of day 12

src/day12.ml

@@ -1,7 +1,5 @@
 open Containers
 
-let infinity = 1000000000
-
 type node =
   { char : char
   ; mutable visited : bool
@@ -10,25 +8,20 @@ type node =
 [@@deriving show]
 
 let parse_grid start_anywhere =
-  let s = ref @@ Vec2.origin in
-  let e = ref @@ Vec2.origin in
-  let unvisited = ref [] in
+  let start_nodes = ref [] in
+  let end_node = ref @@ Vec2.origin in
   let grid =
     Utils.lines_of_input 12
-    |> Grid.of_lines (fun char -> { char; visited = false; distance = infinity })
+    |> Grid.of_lines (fun char -> { char; visited = false; distance = Int.max_int })
   in
   Grid.iter grid (fun pos ->
-    let n = Grid.at_e grid pos in
-    match n.char with
-    | 'E' -> e := pos
-    | 'S' ->
-      n.distance <- 0;
-      s := pos
-    | 'a' when start_anywhere ->
-      n.distance <- 0;
-      unvisited := pos :: !unvisited
+    let node = Grid.at_e grid pos in
+    match node.char with
+    | 'E' -> end_node := pos
+    | 'S' -> start_nodes := pos :: !start_nodes
+    | 'a' when start_anywhere -> start_nodes := pos :: !start_nodes
     | _ -> ());
-  grid, !e, !s, !unvisited
+  grid, !end_node, !start_nodes
 ;;
 
 let reachable f t =
@@ -41,42 +34,45 @@ let reachable f t =
   Char.compare (height f) (height t) >= -1
 ;;
 
-let rec path grid goal current unvisited =
-  let current_node = Grid.at_e grid current in
+let get_neighbors grid node_pos =
+  let node = Grid.at_e grid node_pos in
   Vec2.directions
-  |> List.map (fun dir -> Vec2.(current + dir))
-  |> List.iter (fun neighbor ->
-    match Grid.at grid neighbor with
-    | Some node when (not node.visited) && reachable current_node.char node.char ->
-      let new_distance = current_node.distance + 1 in
-      if new_distance < node.distance then node.distance <- new_distance;
-      if Option.is_none @@ List.find_opt (fun x -> Vec2.(x = neighbor)) !unvisited
-      then unvisited := neighbor :: !unvisited
-    | _ -> ());
-  current_node.visited <- true;
-  unvisited
-    := List.sort
-         (fun a b ->
-           let a = Grid.at_e grid a in
-           let b = Grid.at_e grid b in
-           a.distance - b.distance)
-         !unvisited;
-  match !unvisited with
-  | shortest :: _ when Vec2.(shortest = goal) -> Grid.at grid shortest
-  | shortest :: rest -> path grid goal shortest @@ ref rest
-  | [] -> None
+  |> List.filter_map (fun dir ->
+    let neighbor_pos = Vec2.(node_pos + dir) in
+    match Grid.at grid neighbor_pos with
+    | Some neighbor when reachable node.char neighbor.char -> Some neighbor_pos
+    | _ -> None)
+;;
+
+let path_distance grid goal start_nodes =
+  let set_info position (dist, visited) =
+    let node = Grid.at_e grid position in
+    node.distance <- dist;
+    node.visited <- visited
+  in
+  let get_info position =
+    let { char = _; distance; visited } = Grid.at_e grid position in
+    distance, visited
+  in
+  Utils.djikstra
+    ~set_info
+    ~get_info
+    ~eq:Vec2.( = )
+    ~get_weight:(fun _ -> 1)
+    ~get_neighbors:(get_neighbors grid)
+    ~goal
+    ~start_nodes
+  |> Option.get_exn_or "couldn't find"
 ;;
 
 let%expect_test "Day 12.1" =
-  let grid, e, s, unvisited = parse_grid false in
-  let solution = Option.get_exn_or "coulnd't find" @@ path grid e s @@ ref unvisited in
-  Printf.printf "%d" solution.distance;
+  let grid, goal, start_nodes = parse_grid false in
+  Printf.printf "%i" @@ path_distance grid goal start_nodes;
   [%expect {| 497 |}]
 ;;
 
 let%expect_test "Day 12.2" =
-  let grid, e, s, unvisited = parse_grid true in
-  let solution = Option.get_exn_or "coulnd't find" @@ path grid e s @@ ref unvisited in
-  Printf.printf "%d" solution.distance;
+  let grid, goal, start_nodes = parse_grid true in
+  Printf.printf "%i" @@ path_distance grid goal start_nodes;
   [%expect {| 492 |}]
 ;;

src/utils.ml

@@ -44,6 +44,61 @@ let%expect_test "fib memoized" =
   [%expect {| 1298777728820984005 |}]
 ;;
 
+let djikstra ~eq ~get_info ~set_info ~get_weight ~get_neighbors ~goal ~start_nodes =
+  let seen = ref [] in
+  let rec djikstra' c_node =
+    let c_dist, _ = get_info c_node in
+    set_info c_node (c_dist, true);
+    get_neighbors c_node
+    |> List.iter (fun neighbor ->
+      let n_dist, n_visited = get_info neighbor in
+      if not n_visited
+      then (
+        let new_dist = c_dist + get_weight neighbor in
+        if new_dist < n_dist then set_info neighbor (new_dist, n_visited);
+        seen := List.add_nodup ~eq neighbor !seen));
+    seen
+    := List.fast_sort
+         (fun a b ->
+           let a_dist, _ = get_info a in
+           let b_dist, _ = get_info b in
+           Int.compare a_dist b_dist)
+         !seen;
+    match !seen with
+    | next :: _ when eq next goal ->
+      let dist, _ = get_info next in
+      Some dist
+    | next :: rest ->
+      seen := rest;
+      djikstra' next
+    | [] -> None
+  in
+  List.iter (fun node -> set_info node (0, false)) start_nodes;
+  match start_nodes with
+  | first :: rest ->
+    seen := rest;
+    djikstra' first
+  | _ -> failwith "start_nodes cannot be empty"
+;;
+
+(* not used anywhere so far, but cool as a reference for
+   using locally abstract types!
+*)
+let djikstra_hash (type k) ?(hash = Hashtbl.hash) ?(eq = Stdlib.( = )) slots =
+  let module Tbl =
+    CCHashtbl.Make (struct
+      type t = k
+
+      let equal = eq
+      let hash = hash
+    end)
+  in
+  let info_tbl = Tbl.create slots in
+  let get_info node = Tbl.get_or ~default:(Int.max_int, false) info_tbl node in
+  let set_info node info = Tbl.replace info_tbl node info in
+  djikstra ~eq ~get_info ~set_info
+;;
+
 module Parse = struct
   include Angstrom