#!nix-shell --pure -i ocaml -p ocaml
*)
module CoordSet = Set.Make (struct
type t = int * int
let compare = compare
end)
type direction = North | South | East | West
module CoordMap = Map.Make (struct
type t = (int * int) * direction
let compare = compare
end)
let walls, start_point, end_point =
let with_wall wall (walls, start_point, end_point) =
(CoordSet.add wall walls, start_point, end_point)
in
let with_start_point start_point (walls, _, end_point) =
(walls, Some start_point, end_point)
in
let with_end_point end_point (walls, start_point, _) =
(walls, start_point, Some end_point)
in
let rec parse_line x y = function
| '#' :: xs ->
parse_line (x + 1) y xs |> with_wall (x, y)
| 'E' :: xs ->
parse_line (x + 1) y xs |> with_end_point (x, y)
| 'S' :: xs ->
parse_line (x + 1) y xs |> with_start_point (x, y)
| '.' :: xs ->
parse_line (x + 1) y xs
| x :: xs ->
failwith (Printf.sprintf "invalid character in input %c" x)
| [] ->
(CoordSet.empty, None, None)
in
let rec aux walls start_point end_point y =
try
let new_walls, new_start_point, new_end_point =
read_line () |> String.to_seq |> List.of_seq |> parse_line 0 y
in
let start_point =
match start_point with Some _ as x -> x | None -> new_start_point
in
let end_point =
match end_point with Some _ as x -> x | None -> new_end_point
in
aux (CoordSet.union walls new_walls) start_point end_point (y + 1)
with End_of_file -> (walls, Option.get start_point, Option.get end_point)
in
aux CoordSet.empty None None 0
let next_coords (x, y) = function
| North ->
(x, y - 1)
| South ->
(x, y + 1)
| West ->
(x - 1, y)
| East ->
(x + 1, y)
let invert_direction = function
| North ->
South
| South ->
North
| East ->
West
| West ->
East
let directions = [North; South; East; West]
type state = {dist: int CoordMap.t; queue: (int * (int * int) * direction) list}
let traverse start_point origins =
let initial_state =
{ dist=
List.fold_left
(fun acc origin -> CoordMap.add (start_point, origin) 0 acc)
CoordMap.empty origins
; queue= List.map (fun origin -> (0, start_point, origin)) origins }
in
let is_better_path dist (x, y) dir score =
match CoordMap.find_opt ((x, y), dir) dist with
| Some v ->
v > score
| None ->
true
in
let add_to_queue queue entry = List.sort compare (entry :: queue) in
let process_direction state (score, (x, y), dir) next_dir =
let new_score = score + 1000 in
if not (is_better_path state.dist (x, y) next_dir new_score) then state
else
{ dist= CoordMap.add ((x, y), next_dir) new_score state.dist
; queue= add_to_queue state.queue (new_score, (x, y), next_dir) }
in
let process_forward state (score, (x, y), dir) =
let next_pos = next_coords (x, y) dir in
let new_score = score + 1 in
if
CoordSet.mem next_pos walls
|| not (is_better_path state.dist next_pos dir new_score)
then state
else
{ dist= CoordMap.add (next_pos, dir) new_score state.dist
; queue= add_to_queue state.queue (new_score, next_pos, dir) }
in
let process_node state ((score, pos, dir) as node) =
if is_better_path state.dist pos dir score then state
else
let state_after_directions =
List.fold_left
(fun acc_state next_dir ->
if next_dir = dir then acc_state
else process_direction acc_state node next_dir )
state directions
in
process_forward state_after_directions node
in
let rec process_queue state =
match state.queue with
| [] ->
state.dist
| node :: rest ->
let state' = {state with queue= rest} in
process_queue (process_node state' node)
in
process_queue initial_state
let part1 =
traverse start_point [East]
|> CoordMap.to_list
|> List.filter (fun ((pos, _), _) -> pos = end_point)
|> List.map snd |> List.fold_left min max_int
let _ = part1 |> print_int |> print_newline
let part2 =
let zip o1 o2 =
match (o1, o2) with Some v1, Some v2 -> Some (v1, v2) | _ -> None
in
let from_start = traverse start_point [East] in
let from_end = traverse end_point directions in
let max_x =
walls |> CoordSet.to_list |> List.map fst |> List.fold_left max min_int
in
let max_y =
walls |> CoordSet.to_list |> List.map snd |> List.fold_left max min_int
in
let coordinates =
List.init (max_y + 1) (fun y -> List.init (max_x + 1) (fun x -> (x, y)))
|> List.concat
in
let is_valid_point (x, y) =
directions
|> List.filter_map (fun dir ->
zip
(CoordMap.find_opt ((x, y), dir) from_start)
(CoordMap.find_opt ((x, y), invert_direction dir) from_end) )
|> List.exists (fun (a, b) -> a + b = part1)
in
coordinates |> List.filter is_valid_point
|> List.fold_left (fun acc coord -> CoordSet.add coord acc) CoordSet.empty
|> CoordSet.cardinal |> print_int |> print_newline