134 lines
4.8 KiB
OCaml
134 lines
4.8 KiB
OCaml
(** [pairs_of_ints lst] returns a pair from a list of two elements. *)
|
|
let pairs_of_ints = function
|
|
| [ h; h' ] -> (h, h')
|
|
| _ -> raise (Invalid_argument "pairs_of_ints")
|
|
|
|
(** [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
|
|
|
|
type 'a grid = { grid : 'a array; width : int }
|
|
|
|
(** [grid_is_valid_pos grid (x, y)] returns true if (x, y) is a valid position
|
|
*)
|
|
let grid_is_valid_pos grid (x, y) =
|
|
x >= 0 && x < grid.width && y >= 0 && y < grid.width
|
|
|
|
(** Get the index into the grid from an x, y position. *)
|
|
let grid_idx_by_pos grid (x, y) = x + (y * grid.width)
|
|
|
|
(** Set the value of the position (x, y) to v in grid. *)
|
|
let grid_set_by_pos grid p v =
|
|
assert (grid_is_valid_pos grid p);
|
|
let idx = grid_idx_by_pos grid p in
|
|
grid.grid.(idx) <- v
|
|
|
|
(** Get the value of the position (x, y) in grid. *)
|
|
let grid_get_by_pos grid p =
|
|
assert (grid_is_valid_pos grid p);
|
|
let idx = grid_idx_by_pos grid p in
|
|
grid.grid.(idx)
|
|
|
|
(** [grid_of_rocks w rocks] returns a [w * w] grid with [grid.(x + y * w)]
|
|
indicating whether the space is empty ([=max_int]) or which rock it is (0
|
|
based). *)
|
|
let grid_of_rocks width rocks =
|
|
let grid = { grid = Array.make (width * width) max_int; width } in
|
|
let add_rock idx p = grid_set_by_pos grid p idx in
|
|
List.iteri add_rock rocks;
|
|
grid
|
|
|
|
(** [visit grid has_visited count cost pos] visits the location pos marking it
|
|
as visited and returning a list of [(cost, pos)] pairs of next locations to
|
|
examine.
|
|
|
|
[grid] is the grid of rocks, [has_visited] is an array of bools indicating
|
|
whether a position has already been visited, and [count] is how many rocks
|
|
have fallen. *)
|
|
let visit grid has_visited count cost state =
|
|
if not (grid_is_valid_pos grid state) then []
|
|
else if has_visited.(grid_idx_by_pos grid state) then []
|
|
else if grid_get_by_pos grid state < count then []
|
|
else
|
|
let x, y = state in
|
|
has_visited.(grid_idx_by_pos grid state) <- true;
|
|
[
|
|
(cost + 1, (x + 1, y));
|
|
(cost + 1, (x - 1, y));
|
|
(cost + 1, (x, y + 1));
|
|
(cost + 1, (x, y - 1));
|
|
]
|
|
|
|
(** [grid_of_file w fname] returns a grid of width & height [w] populated with
|
|
rocks described in the file [fname]. *)
|
|
let grid_of_file width fname =
|
|
Aoc.strings_of_file fname
|
|
|> List.map (Aoc.ints_of_string ~sep:",")
|
|
|> List.map pairs_of_ints |> grid_of_rocks width
|
|
|
|
(** [find_route_length count grid] calculates the route from the top-left
|
|
position in [grid] to the bottom right if [count] rocks have fallen. It
|
|
returns [None] if no route is possible or [Some (cost, pos)] if the route is
|
|
possible. *)
|
|
let find_route_length count grid =
|
|
let has_visited = Array.make (Array.length grid.grid) false in
|
|
dijkstra
|
|
(visit grid has_visited count)
|
|
(( = ) (grid.width - 1, grid.width - 1))
|
|
[ (0, (0, 0)) ]
|
|
|
|
(** [part1 count rocks] returns how long it takes to navigate the grid [rocks]
|
|
when [count] rocks have fallen. *)
|
|
let part1 count rocks =
|
|
match find_route_length count rocks with
|
|
| None -> failwith "part1"
|
|
| Some (cost, _) -> cost
|
|
|
|
(** [part2 start_count grid] returns the location of the first rock to fall into
|
|
[grid] which makes it impossible to get from the top-left to bottom-right.
|
|
*)
|
|
let part2 start_count grid =
|
|
let rec impl count =
|
|
match find_route_length count grid with
|
|
| None -> count
|
|
| Some _ -> impl (count + 1)
|
|
in
|
|
let count = impl start_count in
|
|
match Array.find_index (( = ) (count - 1)) grid.grid with
|
|
| None -> failwith "part2"
|
|
| Some idx -> idx
|
|
|
|
(** [string_of_idx width idx] prints the (x, y) location for a given index in a
|
|
grid. *)
|
|
let string_of_idx width idx =
|
|
Printf.sprintf "%d,%d" (idx mod width) (idx / width)
|
|
|
|
(** Width of grid *)
|
|
let width = 71
|
|
|
|
let _ =
|
|
Aoc.main (grid_of_file width)
|
|
[ (string_of_int, part1 1024); (string_of_idx width, part2 1024) ]
|