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

src/day16.ml

@@ -22,12 +22,6 @@ type room =
   }
 [@@deriving show { with_path = false }]
 
-type valve =
-  { flow : int
-  ; valves : (string * int) list
-  }
-[@@deriving show { with_path = false }]
-
 let parse_valves lines =
   let rooms = Hashtbl.create 100 in
   lines
@@ -40,10 +34,10 @@ let parse_valves lines =
   let valves =
     Hashtbl.to_list rooms
     |> List.filter_map (fun (_, { name; flow; tunnels = _ }) ->
-      if flow > 0 || String.(name = "AA") then Some (name, flow) else None)
+      if flow > 0 then Some (name, flow) else None)
   in
   let distances = Hashtbl.create 100 in
-  valves
+  valves @ [ "AA", 0 ]
   |> List.iter (fun (from, _) ->
     let dist_of =
       Fun.(
@@ -62,66 +56,57 @@ let parse_valves lines =
   valves, get_dist
 ;;
 
-(* let get_valve valves name = Option.get_exn_or "poo" @@ Hashtbl.get valves name *)
+let traverse get_dist =
-
+  let rec traverse' valves time elephant current =
-(* (\* I think this algorithm fails for 15.1 example and 15.2 because it can't do *)
+    let options =
-(*    backtracking. But if I dont filter out visited nodes, it takes forever to *)
+      valves
-(*    run. prob need to find the shortest path between each valve thats worth *)
+      |> List.filter_map (fun (name, flow) ->
-(*    opening before doing the traversal. *\) *)
+        let dist = get_dist current name in
-(* let traverse valves = *)
+        let delta = dist + 1 in
-(*   Utils.memo (fun self released current time opened visited -> *)
+        let score = flow * (time - delta) in
-(*     if time = 0 *)
+        if score < 0
-(*     then released, opened *)
+        then None
-(*     else ( *)
+        else (
-(*       let open_current = *)
+          let valves' = List.remove_assoc ~eq:String.equal name valves in
-(*         if current.flow > 0 && (not @@ List.mem current.name opened) *)
+          Some (score + traverse' valves' (time - delta) elephant name)))
-(*         then ( *)
+    in
-(*           let time' = time - 1 in *)
+    (options @ if elephant then [ traverse' valves 26 false "AA" ] else [ 0 ])
-(*           let released' = released + (current.flow * time') in *)
+    |> List.reduce Int.max
-(*           [ self released' current time' (current.name :: opened) visited ]) *)
+    |> Option.get_or ~default:0
-(*         else [] *)
+  in
-(*       in *)
+  traverse'
-(*       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 solve_part_1 lines = *)
-(*   let valves, distances = parse_valves lines in *)
-(*   let start = get_valve valves "AA" in *)
-(*   let released, _ = traverse valves 0 start 30 [] [] in *)
-(*   released *)
-(* ;; *)
-
-(* let%expect_test "Day 16.1 example" = *)
-(*   Printf.printf "%i" @@ solve_part_1 @@ example_lines; *)
-(*   [%expect {| 1358 |}] *)
-(* ;; *)
-
-let%expect_test "Day 16.1" =
-  let _, get_dist = parse_valves example_lines in
-  Printf.printf "%i" @@ get_dist "AA" "HH";
-  [%expect {| 5 |}]
 ;;
 
-(* let solve_part_2 lines = *)
+let%expect_test "example" =
-(*   let valves = parse_valves lines in *)
+  let valves, get_dist = parse_valves example_lines in
-(*   let start = get_valve valves "AA" in *)
+  Printf.printf "part 1: %i\n" @@ traverse get_dist valves 30 false "AA";
-(*   let released, opened = traverse valves 0 start 26 [] [] in *)
+  Printf.printf "part 2: %i\n" @@ traverse get_dist valves 26 true "AA";
-(*   let released', _ = traverse valves released start 26 opened [] in *)
+  [%expect {|
-(*   released' *)
+    part 1: 1651
-(* ;; *)
+    part 2: 1707 |}]
+;;
 
+let%expect_test "Day 16.1 example" =
+  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" =
+  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 |}]
+;;
+
+(* took 106.465 sec >:( *)
 (* let%expect_test "Day 16.2" = *)
-(*   Printf.printf "%i" @@ solve_part_2 @@ Utils.lines_of_input 16; *)
+(*   let valves, get_dist = parse_valves @@ Utils.lines_of_input 16 in *)
-(*   [%expect {| 2353 |}] *)
+(*   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 ->