88 lines
2.9 KiB
OCaml
88 lines
2.9 KiB
OCaml
(** [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 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)
|
|
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 map = IntMap.filter (fun k _ -> k >= 100) map in
|
|
IntMap.fold (fun _ v a -> v + a) map 0
|
|
|
|
let find_start grid =
|
|
match Aoc.Grid.idx_from_opt grid 0 'S' with
|
|
| None -> failwith "find_start"
|
|
| Some x -> Aoc.Grid.pos_of_idx grid x
|
|
|
|
let data_of_file fname =
|
|
let grid = Aoc.Grid.of_file fname in
|
|
let start = find_start grid in
|
|
(grid, start)
|
|
|
|
let _ = Aoc.main data_of_file [ (string_of_int, part1) ]
|