finally solved day 16 part 2 (in 100 seconds :P)

now i should be able to rewrite traverse in a way that doesn't suck

src/day12.ml

@@ -58,7 +58,7 @@ let path_distance grid goal start_nodes =
     ~set_info
     ~get_info
     ~eq:Vec2.( = )
-    ~get_weight:(fun _ -> 1)
+    ~get_weight:(fun _ _ -> 1)
     ~get_neighbors:(get_neighbors grid)
     ~goal
     ~start_nodes

src/day16.ml

@@ -15,7 +15,7 @@ Valve JJ has flow rate=21; tunnel leads to valve II
 |}
 ;;
 
-type valve =
+type room =
   { name : string
   ; flow : int
   ; tunnels : string list
@@ -23,76 +23,90 @@ type valve =
 [@@deriving show { with_path = false }]
 
 let parse_valves lines =
-  let valves = Hashtbl.create 100 in
+  let rooms = Hashtbl.create 100 in
   lines
   |> List.iter (fun line ->
     let int_matches = List.map int_of_string Re.(matches (Pcre.regexp {|\d+|}) line) in
     let valve_matches = Re.(matches (Pcre.regexp {|[A-Z]{2}|}) line) in
     match int_matches, valve_matches with
-    | [ flow ], name :: tunnels -> Hashtbl.add valves name { name; flow; tunnels }
+    | [ flow ], name :: tunnels -> Hashtbl.add rooms name { name; flow; tunnels }
     | _ -> failwith "parse error");
-  valves
+  let valves =
+    Hashtbl.to_list rooms
+    |> List.filter_map (fun (_, { name; flow; tunnels = _ }) ->
+      if flow > 0 then Some (name, flow) else None)
+  in
+  let distances = Hashtbl.create 100 in
+  valves @ [ "AA", 0 ]
+  |> List.iter (fun (from, _) ->
+    let dist_of =
+      Fun.(
+        Utils.djikstra_hash
+          ~get_weight:(fun _ _ -> 1)
+          ~get_neighbors:(fun room ->
+            (Option.get_exn_or "sad" @@ Hashtbl.get rooms room).tunnels)
+          ~start_nodes:[ from ]
+        %> Option.get_exn_or "big sad")
+    in
+    valves
+    |> List.remove_assoc ~eq:String.equal from
+    |> List.iter (fun (valve, _) ->
+      Hashtbl.replace distances (from, valve) (dist_of valve)));
+  let get_dist f t = Option.get_exn_or "huge sad" @@ Hashtbl.get distances (f, t) in
+  valves, get_dist
 ;;
 
-let get_valve valves name = Option.get_exn_or "poo" @@ Hashtbl.get valves name
-
-(* I think this algorithm fails for 15.1 example and 15.2 because it can't do
-   backtracking. But if I dont filter out visited nodes, it takes forever to
-   run. prob need to find the shortest path between each valve thats worth
-   opening before doing the traversal. *)
-let traverse valves =
-  Utils.memo (fun self released current time opened visited ->
-    if time = 0
-    then released, opened
-    else (
-      let open_current =
-        if current.flow > 0 && (not @@ List.mem current.name opened)
-        then (
-          let time' = time - 1 in
-          let released' = released + (current.flow * time') in
-          [ self released' current time' (current.name :: opened) visited ])
-        else []
-      in
-      let follow_a_tunnel =
-        current.tunnels
-        |> List.filter (fun tunnel -> not @@ List.mem tunnel visited)
-        |> List.map (fun name ->
-          let current' = get_valve valves name in
-          let time' = time - 1 in
-          self released current' time' opened (current.name :: visited))
-      in
-      Option.get_or ~default:(released, opened)
-      @@ List.reduce
-           (fun ((a, _) as a') ((b, _) as b') -> if a > b then a' else b')
-           (open_current @ follow_a_tunnel)))
+let traverse get_dist =
+  let rec traverse' valves time elephant current =
+    let options =
+      valves
+      |> List.filter_map (fun (name, flow) ->
+        let dist = get_dist current name in
+        let delta = dist + 1 in
+        let score = flow * (time - delta) in
+        if score < 0
+        then None
+        else (
+          let valves' = List.remove_assoc ~eq:String.equal name valves in
+          Some (score + traverse' valves' (time - delta) elephant name)))
+    in
+    (options @ if elephant then [ traverse' valves 26 false "AA" ] else [ 0 ])
+    |> List.reduce Int.max
+    |> Option.get_or ~default:0
+  in
+  traverse'
 ;;
 
-let solve_part_1 lines =
-  let valves = parse_valves lines in
-  let start = get_valve valves "AA" in
-  let released, _ = traverse valves 0 start 30 [] [] in
-  released
+let%expect_test "example" =
+  let valves, get_dist = parse_valves example_lines in
+  Printf.printf "part 1: %i\n" @@ traverse get_dist valves 30 false "AA";
+  Printf.printf "part 2: %i\n" @@ traverse get_dist valves 26 true "AA";
+  [%expect {|
+    part 1: 1651
+    part 2: 1707 |}]
 ;;
 
 let%expect_test "Day 16.1 example" =
-  Printf.printf "%i" @@ solve_part_1 @@ example_lines;
-  [%expect {| 1358 |}]
+  let valves, get_dist = parse_valves example_lines in
+  Printf.printf "%i\n" @@ traverse get_dist valves 30 false "AA";
+  [%expect {| 1651 |}]
+;;
+
+let%expect_test "Day 16.2 example" =
+  let valves, get_dist = parse_valves example_lines in
+  Printf.printf "%i\n" @@ traverse get_dist valves 26 true "AA";
+  [%expect {| 1707 |}]
 ;;
 
 let%expect_test "Day 16.1" =
-  Printf.printf "%i" @@ solve_part_1 @@ Utils.lines_of_input 16;
+  let valves, get_dist = parse_valves @@ Utils.lines_of_input 16 in
+  Printf.printf "%i\n" @@ traverse get_dist valves 30 false "AA";
   [%expect {| 2265 |}]
 ;;
 
-let solve_part_2 lines =
-  let valves = parse_valves lines in
-  let start = get_valve valves "AA" in
-  let released, opened = traverse valves 0 start 26 [] [] in
-  let released', _ = traverse valves released start 26 opened [] in
-  released'
-;;
-
-let%expect_test "Day 16.2" =
-  Printf.printf "%i" @@ solve_part_2 @@ Utils.lines_of_input 16;
-  [%expect {| 2353 |}]
-;;
+(* took 106.465 sec >:( *)
+(* let%expect_test "Day 16.2" = *)
+(*   let valves, get_dist = parse_valves @@ Utils.lines_of_input 16 in *)
+(*   Printf.printf "%i\n" @@ traverse get_dist valves 26 true "AA"; *)
+(*   [%expect {| 2811 |}] *)
+(* ;; *)

src/utils.ml

@@ -22,7 +22,7 @@ let lines_of_input day =
     ]} *)
 let memo f =
   let open Hashtbl in
-  let cache = create 1000 in
+  let cache = create 100000 in
   let rec f_mem k =
     try find cache k with
     | Not_found ->
@@ -54,7 +54,7 @@ let djikstra ~eq ~get_info ~set_info ~get_weight ~get_neighbors ~goal ~start_nod
       let n_dist, n_visited = get_info neighbor in
       if not n_visited
       then (
-        let new_dist = c_dist + get_weight neighbor in
+        let new_dist = c_dist + get_weight c_node neighbor in
         if new_dist < n_dist then set_info neighbor (new_dist, n_visited);
         seen := List.add_nodup ~eq neighbor !seen));
     seen
@@ -82,10 +82,15 @@ let djikstra ~eq ~get_info ~set_info ~get_weight ~get_neighbors ~goal ~start_nod
   | _ -> 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 djikstra_hash
+  (type k)
+  ?(hash = Hashtbl.hash)
+  ?(eq = Stdlib.( = ))
+  ?(slots = 100)
+  ~get_weight
+  ~get_neighbors
+  ~start_nodes
+  =
   let module Tbl =
     CCHashtbl.Make (struct
       type t = k
@@ -94,10 +99,14 @@ let djikstra_hash (type k) ?(hash = Hashtbl.hash) ?(eq = Stdlib.( = )) slots =
       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
+  let table = Tbl.create slots in
+  let get_info node = Tbl.get_or ~default:(Int.max_int, false) table node in
+  let set_info node info = Tbl.replace table node info in
+  let dist_of node = Option.(Tbl.get table node >|= fun (dist, _) -> dist) in
+  let _ =
+    djikstra ~eq ~get_info ~set_info ~get_weight ~get_neighbors ~goal:None ~start_nodes
+  in
+  dist_of
 ;;
 
 module Parse = struct