2024 day 20 part 2.

bin/day2420.ml

@@ -1,76 +1,81 @@
-(** [dijkstra visit check_end states] executes Dijkstra's algorithm.
-
-    [visit cost state] is called to visit [state] with [cost]. It should mark
-    [state] as visited, and return a list of [(cost, state)] pairs which contain
-    new states to examine. The returned list should be sorted by [cost].
-
-    [check_end state] should return [true] if and only if [state] is an end
-    state.
-
-    [states] is a list of [(cost, state)] pairs ordered by [cost].
-
-    [dijkstra] returns [None] if no path is found to the destination. It returns
-    [Some (cost, state, remaining_states)] if a route is found. [cost] is the
-    cost of getting to [state]. [remaining_states] is a list of the remaining
-    states which can be passed back to [dijkstra] if we want to find further
-    paths. *)
-let rec dijkstra visit check_end =
-  let compare_costs (lhs, _) (rhs, _) = compare lhs rhs in
-  function
-  | [] -> None
-  | (cost, state) :: t ->
-      if check_end state then Some (cost, state)
-      else
-        let new_states = visit cost state |> List.merge compare_costs t in
-        dijkstra visit check_end new_states
-
-let visit grid has_visited cost ((x, y) as state) =
-  if not (Aoc.Grid.pos_is_valid grid state) then []
-  else if has_visited.(Aoc.Grid.idx_of_pos grid state) then []
-  else begin
-    has_visited.(Aoc.Grid.idx_of_pos grid state) <- true;
-    [
-      (cost + 1, (x - 1, y));
-      (cost + 1, (x + 1, y));
-      (cost + 1, (x, y - 1));
-      (cost + 1, (x, y + 1));
-    ]
-  end
-
-let find_cost grid start =
-  let has_visited =
-    Array.init (Aoc.Grid.length grid) (fun x ->
-        Aoc.Grid.get_by_idx grid x = '#')
-  in
-  match
-    dijkstra (visit grid has_visited)
-      (fun x -> Aoc.Grid.get_by_pos_opt grid x = Some 'E')
-      [ (0, start) ]
-  with
-  | None -> failwith "find_cost"
-  | Some (cost, _) -> cost
-
 module IntMap = Map.Make (Int)
 
 let update_value = function None -> Some 1 | Some x -> Some (x + 1)
 
-let find_cost_reductions grid start =
-  let max_cost = find_cost grid start in
+let populate_grid grid start =
+  let costs = Array.make (Aoc.Grid.length grid) max_int in
+  let rec step acc cost = function
+    | [] -> acc
+    | (x, y) :: t ->
+        if
+          Aoc.Grid.pos_is_valid grid (x, y)
+          && Aoc.Grid.get_by_pos grid (x, y) <> '#'
+          && costs.(Aoc.Grid.idx_of_pos grid (x, y)) = max_int
+        then begin
+          costs.(Aoc.Grid.idx_of_pos grid (x, y)) <- cost;
+          if Aoc.Grid.get_by_pos grid (x, y) = 'E' then step acc cost t
+          else
+            step
+              ((x - 1, y) :: (x + 1, y) :: (x, y - 1) :: (x, y + 1) :: acc)
+              cost t
+        end
+        else step acc cost t
+  in
+  let rec dfs cost lst =
+    let next_step = step [] cost lst in
+    if next_step = [] then costs else dfs (cost + 1) next_step
+  in
+  let costs = dfs 0 [ start ] in
+  costs
+
+let manhattan_distance (x, y) (x', y') = abs (x - x') + abs (y - y')
+
+let within_distance (x, y) distance =
+  let rec impl' acc y' x' =
+    if manhattan_distance (x, y) (x', y') > distance then acc
+    else impl' ((x', y') :: acc) y' (x' + 1)
+  in
+  let rec impl acc y' =
+    if y' - y > distance then acc
+    else impl (impl' acc y' (x - (distance - abs (y - y')))) (y' + 1)
+  in
+  impl [] (y - distance)
+
+let find_cost2 map depth_first grid idx length =
+  let saving idx' =
+    let cost = depth_first.(idx) in
+    let cost' = depth_first.(idx') in
+    let saving =
+      cost' - cost
+      - manhattan_distance
+          (Aoc.Grid.pos_of_idx grid idx)
+          (Aoc.Grid.pos_of_idx grid idx')
+    in
+    saving
+  in
+  let rec impl acc = function
+    | [] -> acc
+    | h :: t -> impl (IntMap.update (saving h) update_value acc) t
+  in
+  within_distance (Aoc.Grid.pos_of_idx grid idx) length
+  |> List.filter (Aoc.Grid.pos_is_valid grid)
+  |> List.map (Aoc.Grid.idx_of_pos grid)
+  |> List.filter (fun x -> Aoc.Grid.get_by_idx grid x <> '#')
+  |> List.filter (fun x -> depth_first.(x) - depth_first.(idx) >= 0)
+  |> impl map
+
+let find_cost_reductions2 grid start cheat_length =
+  let depth_first = populate_grid grid start in
   let rec impl acc idx =
     if idx >= Aoc.Grid.length grid then acc
-    else if Aoc.Grid.get_by_idx grid idx <> '#' then impl acc (idx + 1)
     else
-      let grid' = Aoc.Grid.update_idx grid idx '.' in
-      let new_cost = find_cost grid' start in
-      impl (IntMap.update (max_cost - new_cost) update_value acc) (idx + 1)
+      let acc = find_cost2 acc depth_first grid idx cheat_length in
+      impl acc (idx + 1)
   in
   impl IntMap.empty 0
 
-let print_map map = IntMap.iter (fun k v -> Printf.printf "%d -> %d\n" k v) map
-
-let part1 (grid, start) =
-  let map = find_cost_reductions grid start in
-  print_map map;
+let part2 cheat_length (grid, start) =
+  let map = find_cost_reductions2 grid start cheat_length in
   let map = IntMap.filter (fun k _ -> k >= 100) map in
   IntMap.fold (fun _ v a -> v + a) map 0
 
@@ -84,4 +89,5 @@ let data_of_file fname =
   let start = find_start grid in
   (grid, start)
 
-let _ = Aoc.main data_of_file [ (string_of_int, part1) ]
+let _ =
+  Aoc.main data_of_file [ (string_of_int, part2 2); (string_of_int, part2 20) ]